talos/hack/docgen/main.go

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"

	"gopkg.in/yaml.v3"
	"mvdan.cc/gofumpt/format"
)

var tpl = `// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code generated by hack/docgen tool. DO NOT EDIT.

package {{ .Package }}

import (
	"github.com/siderolabs/talos/pkg/machinery/config/encoder"
)

{{ $tick := "` + "`" + `" -}}

{{ range $struct := .Structs -}}
func ({{ $struct.Name }}) Doc() *encoder.Doc {
	doc := &encoder.Doc{
		Type : "{{ $struct.Name }}",
		Comments: [3]string{ "" /* encoder.HeadComment */, "{{ $struct.Text.Comment }}" /* encoder.LineComment */,  "" /* encoder.FootComment */},
		Description: "{{ $struct.Text.Description }}",
		{{ if $struct.AppearsIn -}}
		AppearsIn: []encoder.Appearance{
		{{ range $value := $struct.AppearsIn -}}
			{
				TypeName: "{{ $value.Struct.Name }}",
				FieldName: "{{ $value.FieldName }}",
			},
		{{ end -}}
		},
		{{ end -}}
		Fields: []encoder.Doc{
			{{ range $index, $field := $struct.Fields -}}
				{{ if $field.Tag -}}
					{
					  Name: "{{ $field.Tag }}",
					  Type: "{{ $field.Type }}",
					  Note: "{{ $field.Note }}",
					  Description: "{{ $field.Text.Description }}",
					  Comments: [3]string{ "" /* encoder.HeadComment */, "{{ $field.Text.Comment }}" /* encoder.LineComment */,  "" /* encoder.FootComment */},
	  				  {{ if $field.Text.Values -}}
					  Values : []string{
						{{ range $value := $field.Text.Values -}}
							"{{ $value }}",
						{{ end -}}
					  },
					  {{ end -}}
					},
				{{ else -}}
					{},
				{{- end }}
			{{- end }}

		},
	}

	{{ range $example := $struct.Text.Examples }}
	{{ if $example.Value }}
	doc.AddExample("{{ $example.Name }}", {{ $example.Value }})
	{{ end -}}
	{{ end }}

	{{ range $index, $field := $struct.Fields -}}
		{{ if $field.Tag -}}
			{{ if $field.Text.Examples -}}
				{{ range $example := $field.Text.Examples -}}
					{{- if $example.Value }}
						doc.Fields[{{ $index }}].AddExample("{{ $example.Name }}", {{ $example.Value }})
					{{- end }}
				{{- end }}
			{{- end }}
		{{- end }}
	{{- end }}


	return doc
}
{{ end -}}

// Get{{ .Name }}Doc returns documentation for the file {{ .File }}.
func Get{{ .Name }}Doc() *encoder.FileDoc {
	return &encoder.FileDoc{
		Name: "{{ .Name }}",
		Description: "{{ .Header }}",
		Structs: []*encoder.Doc{
			{{ range $struct := .Structs -}}
			{{ $struct.Name }}{}.Doc(),
			{{ end -}}
		},
	}
}
`

type Doc struct {
	Name    string
	Package string
	Title   string
	Header  string
	File    string
	Structs []*Struct
}

type Struct struct {
	Name      string
	Text      *Text
	Fields    []*Field
	AppearsIn []Appearance
}

type Appearance struct {
	Struct    *Struct
	FieldName string
}

type Example struct {
	Name  string `yaml:"name"`
	Value string `yaml:"value"`
}

type Field struct {
	Name    string
	Type    string
	TypeRef string
	Text    *Text
	Tag     string
	Note    string
}

type Text struct {
	Comment     string         `json:"-"`
	Description string         `json:"description"`
	Examples    []*Example     `json:"examples"`
	Values      []string       `json:"values"`
	Schema      *SchemaWrapper `json:"schema"`
}

func in(p string) (string, error) {
	return filepath.Abs(p)
}

func out(p string) (*os.File, error) {
	abs, err := filepath.Abs(p)
	if err != nil {
		return nil, err
	}

	return os.Create(abs)
}

type structType struct {
	name string
	text *Text
	pos  token.Pos
	node *ast.StructType
}

type aliasType struct {
	fieldType    string
	fieldTypeRef string
}

func collectStructs(node ast.Node) ([]*structType, map[string]aliasType) {
	structs := []*structType{}
	aliases := map[string]aliasType{}

	collectStructs := func(n ast.Node) bool {
		g, ok := n.(*ast.GenDecl)
		if !ok {
			return true
		}

		isAlias := false

		if g.Doc != nil {
			for _, comment := range g.Doc.List {
				if strings.Contains(comment.Text, "docgen:nodoc") {
					return true
				}

				if strings.Contains(comment.Text, "docgen:alias") {
					isAlias = true
				}
			}
		}

		for _, spec := range g.Specs {
			t, ok := spec.(*ast.TypeSpec)
			if !ok {
				return true
			}

			if t.Type == nil {
				return true
			}

			x, ok := t.Type.(*ast.StructType)
			if !ok {
				if isAlias {
					aliases[t.Name.Name] = aliasType{
						fieldType:    formatFieldType(t.Type),
						fieldTypeRef: getFieldType(t.Type),
					}
				}

				return true
			}

			structName := t.Name.Name

			text := &Text{}

			if t.Doc != nil {
				text = parseComment([]byte(t.Doc.Text()))
			} else if g.Doc != nil {
				text = parseComment([]byte(g.Doc.Text()))
			}

			s := &structType{
				name: structName,
				text: text,
				node: x,
				pos:  x.Pos(),
			}

			structs = append(structs, s)
		}

		return true
	}

	ast.Inspect(node, collectStructs)

	return structs, aliases
}

func parseComment(comment []byte) *Text {
	text := &Text{}
	if err := yaml.Unmarshal(comment, text); err != nil {
		lines := strings.Split(string(comment), "\n")
		for i := range lines {
			lines[i] = strings.TrimLeft(lines[i], "\t")
		}

		// not yaml, fallback
		text.Description = strings.Join(lines, "\n")
		// take only the first line from the Description for the comment
		text.Comment = lines[0]

		// try to parse everything except for the first line as yaml
		if err = yaml.Unmarshal([]byte(strings.Join(lines[1:], "\n")), text); err == nil {
			// if parsed, remove it from the description
			text.Description = text.Comment
		}
	} else {
		text.Description = strings.TrimSpace(text.Description)
		// take only the first line from the Description for the comment
		text.Comment = strings.Split(text.Description, "\n")[0]
	}

	text.Comment = escape(text.Comment)

	text.Description = escape(text.Description)
	for _, example := range text.Examples {
		example.Name = escape(example.Name)
		example.Value = strings.TrimSpace(example.Value)
	}

	return text
}

func getFieldType(p interface{}) string {
	if m, ok := p.(*ast.MapType); ok {
		return getFieldType(m.Value)
	}

	switch t := p.(type) {
	case *ast.Ident:
		return t.Name
	case *ast.ArrayType:
		return getFieldType(p.(*ast.ArrayType).Elt)
	case *ast.StarExpr:
		return getFieldType(t.X)
	case *ast.SelectorExpr:
		return getFieldType(t.Sel)
	default:
		return ""
	}
}

func formatFieldType(p interface{}) string {
	if m, ok := p.(*ast.MapType); ok {
		return fmt.Sprintf("map[%s]%s", formatFieldType(m.Key), formatFieldType(m.Value))
	}

	switch t := p.(type) {
	case *ast.Ident:
		return t.Name
	case *ast.ArrayType:
		return "[]" + formatFieldType(p.(*ast.ArrayType).Elt)
	case *ast.StructType:
		return "struct"
	case *ast.StarExpr:
		return formatFieldType(t.X)
	case *ast.SelectorExpr:
		return formatFieldType(t.Sel)
	default:
		log.Printf("unknown: %#v", t)

		return ""
	}
}

func escape(value string) string {
	value = strings.ReplaceAll(value, `"`, `\"`)
	value = strings.ReplaceAll(value, "\n", `\n`)

	return strings.TrimSpace(value)
}

func collectFields(s *structType, aliases map[string]aliasType) (fields []*Field) {
	fields = []*Field{}

	for _, f := range s.node.Fields.List {
		if f.Names == nil {
			// This is an embedded struct.
			continue
		}

		name := f.Names[0].Name

		if f.Doc == nil {
			log.Fatalf("field %q is missing a documentation", name)
		}

		if strings.Contains(f.Doc.Text(), "docgen:nodoc") {
			fields = append(fields, &Field{Type: "unknown"})

			continue
		}

		if f.Tag == nil {
			log.Fatalf("field %q is missing a tag", name)
		}

		fieldType := formatFieldType(f.Type)

		if alias, ok := aliases[fieldType]; ok {
			fieldType = alias.fieldType
		}

		fieldTypeRef := getFieldType(f.Type)

		if alias, ok := aliases[fieldTypeRef]; ok {
			fieldTypeRef = alias.fieldTypeRef
		}

		tag := reflect.StructTag(strings.Trim(f.Tag.Value, "`"))
		yamlTag := tag.Get("yaml")
		yamlTag = strings.Split(yamlTag, ",")[0]

		if yamlTag == "" {
			log.Fatalf("field %q is missing a `yaml` tag or name in it", name)
		}

		text := parseComment([]byte(f.Doc.Text()))

		field := &Field{
			Name:    name,
			Tag:     yamlTag,
			Type:    fieldType,
			TypeRef: fieldTypeRef,
			Text:    text,
		}

		if f.Comment != nil {
			field.Note = escape(f.Comment.Text())
		}

		fields = append(fields, field)
	}

	return fields
}

func render(doc *Doc, dest string) {
	t := template.Must(template.New("docfile.tpl").Parse(tpl))
	buf := bytes.Buffer{}

	err := t.Execute(&buf, doc)
	if err != nil {
		log.Fatalf("failed to render template: %v", err)
	}

	formatted, err := format.Source(buf.Bytes(), format.Options{})
	if err != nil {
		log.Printf("data: %s", buf.Bytes())
		log.Fatalf("failed to format source: %v", err)
	}

	out, err := out(dest)
	if err != nil {
		log.Fatalf("failed to create output file: %v", err)
	}
	defer out.Close()

	_, err = out.Write(formatted)

	if err != nil {
		log.Fatalf("failed to write output file: %v", err)
	}
}

func processFile(inputFile, outputFile, schemaOutputFile, versionTagFile, typeName string) {
	abs, err := in(inputFile)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("creating package file set: %q\n", abs)

	fset := token.NewFileSet()

	node, err := parser.ParseFile(fset, abs, nil, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}

	packageName := node.Name.Name

	tokenFile := fset.File(node.Pos())
	if tokenFile == nil {
		log.Fatalf("No token")
	}

	fmt.Printf("parsing file in package %q: %s\n", packageName, tokenFile.Name())

	structs, aliases := collectStructs(node)

	if len(structs) == 0 {
		log.Fatalf("failed to find types that could be documented in %s", abs)
	}

	doc := &Doc{
		Package: packageName,
		Structs: []*Struct{},
	}

	extraExamples := map[string][]*Example{}
	backReferences := map[string][]Appearance{}

	for _, s := range structs {
		fmt.Printf("generating docs for type: %q\n", s.name)

		fields := collectFields(s, aliases)

		s := &Struct{
			Name:   s.name,
			Text:   s.text,
			Fields: fields,
		}

		for _, field := range fields {
			if field.TypeRef == "" {
				continue
			}

			if len(field.Text.Examples) > 0 {
				extraExamples[field.TypeRef] = append(extraExamples[field.TypeRef], field.Text.Examples...)
			}

			backReferences[field.TypeRef] = append(backReferences[field.TypeRef], Appearance{
				Struct:    s,
				FieldName: field.Tag,
			})
		}

		doc.Structs = append(doc.Structs, s)
	}

	for _, s := range doc.Structs {
		if extra, ok := extraExamples[s.Name]; ok {
			s.Text.Examples = append(s.Text.Examples, extra...)
		}

		if ref, ok := backReferences[s.Name]; ok {
			s.AppearsIn = append(s.AppearsIn, ref...)
		}
	}

	if err == nil {
		doc.Package = node.Name.Name
		doc.Name = typeName

		if node.Doc != nil {
			doc.Header = escape(node.Doc.Text())
		}
	}

	doc.File = outputFile
	render(doc, outputFile)

	if schemaOutputFile != "" {
		renderSchema(doc, schemaOutputFile, versionTagFile)
	}
}

func main() {
	flag.Parse()

	if flag.NArg() != 3 && flag.NArg() != 5 {
		log.Fatalf("unexpected number of args: %d", flag.NArg())
	}

	inputFile := flag.Arg(0)
	outputFile := flag.Arg(1)
	typeName := flag.Arg(2)
	jsonSchemaOutputFile := flag.Arg(3)
	versionTagFile := flag.Arg(4)

	processFile(inputFile, outputFile, jsonSchemaOutputFile, versionTagFile, typeName)
}