github.com/oam-dev/kubevela@v1.9.11/references/cuegen/generator.go (about)

     1  /*
     2  Copyright 2023 The KubeVela Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package cuegen
    18  
    19  import (
    20  	"fmt"
    21  	goast "go/ast"
    22  	gotypes "go/types"
    23  	"io"
    24  	"strings"
    25  
    26  	cueast "cuelang.org/go/cue/ast"
    27  	"cuelang.org/go/cue/ast/astutil"
    28  	cueformat "cuelang.org/go/cue/format"
    29  	"golang.org/x/tools/go/packages"
    30  )
    31  
    32  // Generator generates CUE schema from Go struct.
    33  type Generator struct {
    34  	// immutable
    35  	pkg   *packages.Package
    36  	types typeInfo
    37  
    38  	opts *options
    39  }
    40  
    41  // NewGenerator creates a new generator with given file or package path.
    42  func NewGenerator(f string) (*Generator, error) {
    43  	pkg, err := loadPackage(f)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	types := getTypeInfo(pkg)
    49  
    50  	g := &Generator{
    51  		pkg:   pkg,
    52  		types: types,
    53  		opts:  newDefaultOptions(),
    54  	}
    55  
    56  	return g, nil
    57  }
    58  
    59  // Package returns internal package struct, which should be read-only.
    60  func (g *Generator) Package() *packages.Package {
    61  	return g.pkg
    62  }
    63  
    64  // Generate generates CUE schema from Go struct and writes to w.
    65  // And it can be called multiple times with different options.
    66  //
    67  // NB: it's not thread-safe.
    68  func (g *Generator) Generate(opts ...Option) (decls []Decl, _ error) {
    69  	g.opts = newDefaultOptions() // reset options for each call
    70  	for _, opt := range opts {
    71  		if opt != nil {
    72  			opt(g.opts)
    73  		}
    74  	}
    75  
    76  	for _, syntax := range g.pkg.Syntax {
    77  		for _, decl := range syntax.Decls {
    78  			if d, ok := decl.(*goast.GenDecl); ok {
    79  				t, err := g.convertDecls(d)
    80  				if err != nil {
    81  					return nil, err
    82  				}
    83  				decls = append(decls, t...)
    84  			}
    85  		}
    86  	}
    87  
    88  	return decls, nil
    89  }
    90  
    91  // Format formats CUE ast decls with package header and writes to w.
    92  func (g *Generator) Format(w io.Writer, decls []Decl) error {
    93  	if w == nil {
    94  		return fmt.Errorf("nil writer")
    95  	}
    96  	if len(decls) == 0 {
    97  		return fmt.Errorf("invalid decls")
    98  	}
    99  
   100  	pkg := &cueast.Package{Name: Ident(g.pkg.Name, false)}
   101  
   102  	f := &cueast.File{Decls: []cueast.Decl{pkg}}
   103  	for _, decl := range decls {
   104  		if decl == nil {
   105  			continue
   106  		}
   107  		f.Decls = append(f.Decls, decl.Build())
   108  	}
   109  
   110  	if err := astutil.Sanitize(f); err != nil {
   111  		return err
   112  	}
   113  
   114  	b, err := cueformat.Node(f, cueformat.Simplify())
   115  	if err != nil {
   116  		return err
   117  	}
   118  
   119  	_, err = w.Write(b)
   120  	return err
   121  }
   122  
   123  // loadPackage loads a package from given path.
   124  func loadPackage(p string) (*packages.Package, error) {
   125  	cfg := &packages.Config{
   126  		Mode: packages.NeedName | packages.NeedFiles | packages.NeedCompiledGoFiles |
   127  			packages.NeedImports | packages.NeedTypes | packages.NeedTypesSizes |
   128  			packages.NeedSyntax | packages.NeedTypesInfo | packages.NeedDeps |
   129  			packages.NeedModule,
   130  	}
   131  
   132  	pkgs, err := packages.Load(cfg, []string{p}...)
   133  	if err != nil {
   134  		return nil, err
   135  	}
   136  	if len(pkgs) != 1 {
   137  		return nil, fmt.Errorf("expected one package, got %d", len(pkgs))
   138  	}
   139  
   140  	// only need to check the first package
   141  	pkg := pkgs[0]
   142  	if pkg.Errors != nil {
   143  		errs := make([]string, 0, len(pkg.Errors))
   144  		for _, e := range pkg.Errors {
   145  			errs = append(errs, fmt.Sprintf("\t%s: %v", pkg.PkgPath, e))
   146  		}
   147  		return nil, fmt.Errorf("could not load Go packages:\n%s", strings.Join(errs, "\n"))
   148  	}
   149  
   150  	return pkg, nil
   151  }
   152  
   153  type typeInfo map[gotypes.Type]*goast.StructType
   154  
   155  func getTypeInfo(p *packages.Package) typeInfo {
   156  	m := make(typeInfo)
   157  
   158  	for _, f := range p.Syntax {
   159  		goast.Inspect(f, func(n goast.Node) bool {
   160  			// record all struct types
   161  			if t, ok := n.(*goast.StructType); ok {
   162  				m[p.TypesInfo.TypeOf(t)] = t
   163  			}
   164  			return true
   165  		})
   166  	}
   167  
   168  	return m
   169  }