github.com/linchen2chris/hugo@v0.0.0-20230307053224-cec209389705/markup/goldmark/toc.go (about)

     1  // Copyright 2019 The Hugo Authors. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  // http://www.apache.org/licenses/LICENSE-2.0
     7  //
     8  // Unless required by applicable law or agreed to in writing, software
     9  // distributed under the License is distributed on an "AS IS" BASIS,
    10  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package goldmark
    15  
    16  import (
    17  	"bytes"
    18  
    19  	"github.com/gohugoio/hugo/markup/tableofcontents"
    20  
    21  	"github.com/yuin/goldmark"
    22  	"github.com/yuin/goldmark/ast"
    23  	"github.com/yuin/goldmark/parser"
    24  	"github.com/yuin/goldmark/renderer"
    25  	"github.com/yuin/goldmark/text"
    26  	"github.com/yuin/goldmark/util"
    27  )
    28  
    29  var (
    30  	tocResultKey = parser.NewContextKey()
    31  	tocEnableKey = parser.NewContextKey()
    32  )
    33  
    34  type tocTransformer struct {
    35  	r renderer.Renderer
    36  }
    37  
    38  func (t *tocTransformer) Transform(n *ast.Document, reader text.Reader, pc parser.Context) {
    39  	if b, ok := pc.Get(tocEnableKey).(bool); !ok || !b {
    40  		return
    41  	}
    42  
    43  	var (
    44  		toc         tableofcontents.Builder
    45  		tocHeading  = &tableofcontents.Heading{}
    46  		level       int
    47  		row         = -1
    48  		inHeading   bool
    49  		headingText bytes.Buffer
    50  	)
    51  
    52  	ast.Walk(n, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
    53  		s := ast.WalkStatus(ast.WalkContinue)
    54  		if n.Kind() == ast.KindHeading {
    55  			if inHeading && !entering {
    56  				tocHeading.Title = headingText.String()
    57  				headingText.Reset()
    58  				toc.AddAt(tocHeading, row, level-1)
    59  				tocHeading = &tableofcontents.Heading{}
    60  				inHeading = false
    61  				return s, nil
    62  			}
    63  
    64  			inHeading = true
    65  		}
    66  
    67  		if !(inHeading && entering) {
    68  			return s, nil
    69  		}
    70  
    71  		switch n.Kind() {
    72  		case ast.KindHeading:
    73  			heading := n.(*ast.Heading)
    74  			level = heading.Level
    75  
    76  			if level == 1 || row == -1 {
    77  				row++
    78  			}
    79  
    80  			id, found := heading.AttributeString("id")
    81  			if found {
    82  				tocHeading.ID = string(id.([]byte))
    83  			}
    84  		case
    85  			ast.KindCodeSpan,
    86  			ast.KindLink,
    87  			ast.KindImage,
    88  			ast.KindEmphasis:
    89  			err := t.r.Render(&headingText, reader.Source(), n)
    90  			if err != nil {
    91  				return s, err
    92  			}
    93  
    94  			return ast.WalkSkipChildren, nil
    95  		case
    96  			ast.KindAutoLink,
    97  			ast.KindRawHTML,
    98  			ast.KindText,
    99  			ast.KindString:
   100  			err := t.r.Render(&headingText, reader.Source(), n)
   101  			if err != nil {
   102  				return s, err
   103  			}
   104  		}
   105  
   106  		return s, nil
   107  	})
   108  
   109  	pc.Set(tocResultKey, toc.Build())
   110  }
   111  
   112  type tocExtension struct {
   113  	options []renderer.Option
   114  }
   115  
   116  func newTocExtension(options []renderer.Option) goldmark.Extender {
   117  	return &tocExtension{
   118  		options: options,
   119  	}
   120  }
   121  
   122  func (e *tocExtension) Extend(m goldmark.Markdown) {
   123  	r := goldmark.DefaultRenderer()
   124  	r.AddOptions(e.options...)
   125  	m.Parser().AddOptions(parser.WithASTTransformers(util.Prioritized(&tocTransformer{
   126  		r: r,
   127  	}, 10)))
   128  }