go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/starlark/docgen/generator.go

// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package docgen generates documentation from Starlark code.
package docgen

import (
	"bytes"
	"fmt"
	"regexp"
	"strings"
	"text/template"

	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/starlark/docgen/ast"
	"go.chromium.org/luci/starlark/docgen/docstring"
	"go.chromium.org/luci/starlark/docgen/symbols"
)

// Generator renders text templates that have access to parsed structured
// representation of Starlark modules.
//
// The templates use them to inject documentation extracted from Starlark into
// appropriate places.
//
// It is a cache of the current loaded modules, which enables following
// Render() calls to be much faster.
type Generator struct {
	// Normalize normalizes a load() reference relative to the context of the
	// current starlark file.
	Normalize func(parent, ref string) (string, error)
	// Starlark produces Starlark module's source code.
	//
	// It is then parsed by the generator to extract documentation from it.
	Starlark func(module string) (src string, err error)

	loader *symbols.Loader    // knows how to load symbols from starlark modules
	links  map[string]*symbol // full name -> symbol we can link to
}

// Render renders the given text template in an environment with access to
// parsed structured Starlark comments.
//
// Loaded modules are kept as a cache in Generator, making the rendering of
// multiple starlark files faster.
func (g *Generator) Render(templ string) ([]byte, error) {
	if g.loader == nil {
		g.loader = &symbols.Loader{Normalize: g.Normalize, Source: g.Starlark}
		g.links = map[string]*symbol{}
	}

	t, err := template.New("main").Funcs(g.funcMap()).Parse(templ)
	if err != nil {
		return nil, err
	}

	buf := bytes.Buffer{}
	if err := t.Execute(&buf, nil); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// funcMap are functions available to templates.
func (g *Generator) funcMap() template.FuncMap {
	return template.FuncMap{
		"EscapeMD":       escapeMD,
		"Symbol":         g.symbol,
		"LinkifySymbols": g.linkifySymbols,
	}
}

// escapeMD makes sure 's' gets rendered as is in markdown.
func escapeMD(s string) string {
	return strings.Replace(s, "*", "\\*", -1)
}

func (g *Generator) load(module string) (*symbols.Struct, error) {
	mod, err := g.loader.Load(module)
	if err != nil {
		return nil, err
	}

	// Transform lucicfg.rule(...) definitions to pick up docstrings and arguments
	// of the rule implementation. We replace `var = lucicfg.rule(impl = f)` with
	// `var = f`.
	return mod.Transform(func(s symbols.Symbol) (symbols.Symbol, error) {
		inv, ok := s.(*symbols.Invocation)
		if !ok {
			return s, nil
		}
		// Rule constructor symbols are marked with RuleCtor tag.
		targetTags := inv.Func().Doc().RemarkBlock("DocTags").Body
		if !strings.Contains(targetTags, "RuleCtor") {
			return s, nil
		}
		// Find a symbol assigned to 'impl' kwarg and return it, so that it is
		// used instead of lucicfg.rule(...) invocation. Give it the name of 's'.
		for _, arg := range inv.Args() {
			if arg.Name() == "impl" {
				return symbols.NewAlias(s.Name(), arg), nil
			}
		}
		return nil, fmt.Errorf("cannot resolve rule constructor call in %s, no `impl` kwarg", s)
	})
}

// symbol returns a symbol from the given module.
//
// lookup is a field path, e.g. "a.b.c". "a" will be searched for in the
// top-level dict of the module. If empty, the module itself will be returned.
//
// If the requested symbol can't be found, returns a broken symbol.
func (g *Generator) symbol(module, lookup string) (*symbol, error) {
	mod, err := g.load(module)
	if err != nil {
		return nil, err
	}

	var lookupPath []string
	if lookup != "" {
		lookupPath = strings.Split(lookup, ".")
	}

	sym := &symbol{
		Symbol:   symbols.Lookup(mod, lookupPath...),
		Module:   module,
		FullName: lookup,
	}

	// Let automatic linkifier know about the loaded symbols so it can start
	// generating links to them if it encounters them in the text.
	syms, _ := sym.Symbols()
	for _, s := range syms {
		g.links[s.FullName] = s
	}

	return sym, nil
}

// This matches a.b.c(...). '(...)' part is important, otherwise there is a ton
// of undesired matches in various code snippets.
var symRefRe = regexp.MustCompile(`\w+(\.\w+)*(\(\.\.\.\))`)

// linkifySymbols replaces recognized symbol names with markdown links to
// symbols.
func (g *Generator) linkifySymbols(text string) string {
	return symRefRe.ReplaceAllStringFunc(text, func(match string) string {
		if sym := g.links[strings.TrimSuffix(match, "(...)")]; sym != nil {
			return fmt.Sprintf("[%s](#%s)", match, sym.Anchor())
		}
		return match
	})
}

////////////////////////////////////////////////////////////////////////////////

// symbol is what we expose to the templates.
//
// It is mostly symbols.Symbol, except we add few useful utility fields and
// methods.
type symbol struct {
	symbols.Symbol

	Module   string // module name used to load this symbol
	FullName string // field path from module's top dict to this symbol

	doc  *docstring.Parsed // lazily redacted doc, see Doc()
	tags stringset.Set     // lazily extracted from "DocTags" remarks section
}

// Flavor returns one of "func", "var", "struct", "unknown".
func (s *symbol) Flavor() string {
	switch s.Symbol.(type) {
	case *symbols.Term:
		switch s.Symbol.Def().(type) {
		case *ast.Function:
			return "func"
		case *ast.Var:
			return "var"
		default:
			return "unknown"
		}
	case *symbols.Invocation:
		return "inv"
	case *symbols.Struct:
		return "struct"
	default:
		return "unknown"
	}
}

// Doc is a parsed docstring for this symbol.
//
// An argument called `ctx` is kicked out since it is part of the internal
// lucicfg API.
func (s *symbol) Doc() *docstring.Parsed {
	if s.doc == nil {
		s.doc = s.Symbol.Doc()
		for i, block := range s.doc.Fields {
			if block.Title == "Args" {
				filtered := block.Fields[:0]
				for _, field := range block.Fields {
					if field.Name != "ctx" {
						filtered = append(filtered, field)
					}
				}
				block.Fields = filtered
				s.doc.Fields[i] = block
				break
			}
		}
	}
	return s.doc
}

// HasDocTag returns true if the docstring has a section "DocTags" and the
// given tag is listed there.
//
// Used to mark some symbols as advanced, or experimental.
func (s *symbol) HasDocTag(tag string) bool {
	if s.tags == nil {
		s.tags = stringset.Set{}
		for _, word := range strings.Fields(s.Doc().RemarkBlock("DocTags").Body) {
			s.tags.Add(strings.ToLower(strings.Trim(word, ".,")))
		}
	}
	return s.tags.Has(strings.ToLower(tag))
}

// Symbols returns nested symbols.
//
// If `flavors` is not empty, it specifies what kinds of symbols to keep.
// Possible variants: "func", "var", "inv", "struct".
func (s *symbol) Symbols(flavors ...string) (out []*symbol, err error) {
	strct, _ := s.Symbol.(*symbols.Struct)
	if strct == nil {
		return nil, fmt.Errorf("%q is not a struct", s.FullName)
	}

	keepFlavors := stringset.NewFromSlice(flavors...)

	for _, sym := range strct.Symbols() {
		fullName := ""
		if s.FullName != "" {
			fullName = s.FullName + "." + sym.Name()
		} else {
			fullName = sym.Name()
		}

		sym := &symbol{
			Symbol:   sym,
			Module:   s.Module,
			FullName: fullName,
		}

		if keepFlavors.Len() == 0 || keepFlavors.Has(sym.Flavor()) {
			out = append(out, sym)
		}
	}

	return
}

// Anchor returns a markdown anchor name that can be used to link to some part
// of this symbol's documentation from other parts of the doc.
func (s *symbol) Anchor(sub ...string) string {
	// Gitiles markdown doesn't like '_' in explicitly defined anchors for some
	// reason. It also doesn't like runs of hyphens.
	name := strings.ReplaceAll(s.FullName, "_", "-")
	anchor := strings.Join(append([]string{name}, sub...), "-")
	filtered := ""
	last := '\x00'
	for _, ch := range strings.Trim(anchor, "-") {
		if ch == '-' && last == '-' {
			continue
		}
		last = ch
		filtered += string(ch)
	}
	return filtered
}

// InvocationSnippet returns a snippet showing how a function represented by
// this symbol can be called.
//
// Like this:
//
//	luci.recipe(
//	    # Required arguments.
//	    name,
//	    cipd_package,
//
//	    # Optional arguments.
//	    cipd_version = None,
//	    recipe = None,
//
//	    **kwargs,
//	)
//
// This is apparently very non-trivial to generate using text/template while
// keeping all spaces and newlines strict.
func (s *symbol) InvocationSnippet() string {
	var req, opt, variadric []string
	for _, f := range s.Doc().Args() {
		switch {
		case strings.HasPrefix(f.Name, "*"):
			variadric = append(variadric, f.Name)
		case isRequiredField(f):
			req = append(req, f.Name)
		default:
			opt = append(opt, fmt.Sprintf("%s = None", f.Name))
		}
	}

	b := &strings.Builder{}

	writeSection := func(section string, args []string) {
		if len(args) != 0 {
			b.WriteString(section)
			for _, a := range args {
				fmt.Fprintf(b, "    %s,\n", a)
			}
		}
	}

	fmt.Fprintf(b, "%s(", s.FullName)
	if all := append(append(req, opt...), variadric...); len(all) <= 3 {
		// Use a compact form when we have only very few arguments.
		b.WriteString(strings.Join(all, ", "))
	} else {
		writeSection("\n    # Required arguments.\n", req)
		writeSection("\n    # Optional arguments.\n", opt)
		writeSection("\n", variadric)
	}
	fmt.Fprintf(b, ")")
	return b.String()
}

// isRequiredField takes a field description and tries to figure out whether
// this field is required.
//
// Does this by searching for "Required." suffix. Very robust.
func isRequiredField(f docstring.Field) bool {
	return strings.HasSuffix(f.Desc, "Required.")
}