@@ -6,16 +6,15 @@ package yamlgen
66import (
77 "bytes"
88 "context"
9- "fmt"
109 "go/ast"
1110 "go/parser"
1211 "go/token"
1312 "go/types"
14- "io/ioutil"
1513 "os"
1614 "os/exec"
1715 "path/filepath"
1816 "regexp"
17+ "strings"
1918
2019 "github.com/dave/jennifer/jen"
2120 "github.com/pkg/errors"
@@ -31,18 +30,15 @@ func GenGoCode(src []byte) (string, error) {
3130 // Create new main file.
3231 fset := token .NewFileSet ()
3332 generatedCode := jen .NewFile ("main" )
33+ // Don't really need to format here, saves time.
34+ // generatedCode.NoFormat = true
3435
3536 // Parse source file.
3637 f , err := parser .ParseFile (fset , "" , src , parser .AllErrors )
3738 if err != nil {
3839 return "" , err
3940 }
4041
41- // Add imports if needed(will not be used if not required in code).
42- for _ , s := range f .Imports {
43- generatedCode .ImportName (s .Path .Value [1 :len (s .Path .Value )- 1 ], "" )
44- }
45-
4642 // Init statements for structs.
4743 var init []jen.Code
4844 // Declare config map, i.e, `configs := map[string]interface{}{}`.
@@ -54,33 +50,70 @@ func GenGoCode(src []byte) (string, error) {
5450 if genericDecl , ok := decl .(* ast.GenDecl ); ok {
5551 // Check if declaration spec is `type`.
5652 if typeDecl , ok := genericDecl .Specs [0 ].(* ast.TypeSpec ); ok {
57- var structFields []jen.Code
5853 // Cast to `type struct`.
5954 structDecl , ok := typeDecl .Type .(* ast.StructType )
6055 if ! ok {
6156 generatedCode .Type ().Id (typeDecl .Name .Name ).Id (string (src [typeDecl .Type .Pos ()- 1 : typeDecl .Type .End ()- 1 ]))
6257 continue
6358 }
64- fields := structDecl .Fields .List
59+
60+ var structFields []jen.Code
6561 arrayInit := make (jen.Dict )
6662
6763 // Loop and generate fields for each field.
64+ fields := structDecl .Fields .List
6865 for _ , field := range fields {
6966 // Each field might have multiple names.
7067 names := field .Names
7168 for _ , n := range names {
7269 if n .IsExported () {
7370 pos := n .Obj .Decl .(* ast.Field )
7471
75- // Check if field is a slice type.
76- sliceRe := regexp .MustCompile (`.*\[.*\].*` )
77- if sliceRe .MatchString (types .ExprString (field .Type )) {
78- arrayInit [jen .Id (n .Name )] = jen .Id (string (src [pos .Type .Pos ()- 1 : pos .Type .End ()- 1 ])).Values (jen .Id (string (src [pos .Type .Pos ()+ 1 : pos .Type .End ()- 1 ])).Values ())
72+ // Copy struct field to generated code, with imports in case of other package imported field.
73+ if pos .Tag != nil {
74+ typeStr := string (src [pos .Type .Pos ()- 1 : pos .Type .End ()- 1 ])
75+ // Check if field is imported from other package.
76+ if strings .Contains (typeStr , "." ) {
77+ typeArr := strings .SplitN (typeStr , "." , 2 )
78+ // Match the import name with the import statement.
79+ for _ , s := range f .Imports {
80+ moduleName := ""
81+ // Choose to copy same alias as in source file or use package name.
82+ if s .Name == nil {
83+ _ , moduleName = filepath .Split (s .Path .Value [1 : len (s .Path .Value )- 1 ])
84+ generatedCode .ImportName (s .Path .Value [1 :len (s .Path .Value )- 1 ], moduleName )
85+ } else {
86+ moduleName = s .Name .String ()
87+ generatedCode .ImportAlias (s .Path .Value [1 :len (s .Path .Value )- 1 ], moduleName )
88+ }
89+ // Add field to struct only if import name matches.
90+ if moduleName == typeArr [0 ] {
91+ structFields = append (structFields , jen .Id (n .Name ).Qual (s .Path .Value [1 :len (s .Path .Value )- 1 ], typeArr [1 ]).Id (pos .Tag .Value ))
92+ }
93+ }
94+ } else {
95+ structFields = append (structFields , jen .Id (n .Name ).Id (string (src [pos .Type .Pos ()- 1 :pos .Type .End ()- 1 ])).Id (pos .Tag .Value ))
96+ }
7997 }
8098
81- // Copy struct field to generated code.
82- if pos .Tag != nil {
83- structFields = append (structFields , jen .Id (n .Name ).Id (string (src [pos .Type .Pos ()- 1 :pos .Type .End ()- 1 ])).Id (pos .Tag .Value ))
99+ // Check if field is a slice type.
100+ sliceRe := regexp .MustCompile (`^\[.*\].*$` )
101+ typeStr := types .ExprString (field .Type )
102+ if sliceRe .MatchString (typeStr ) {
103+ iArr := "[]int"
104+ fArr := "[]float"
105+ cArr := "[]complex"
106+ uArr := "[]uint"
107+ switch typeStr {
108+ case "[]bool" , "[]string" , "[]byte" , "[]rune" ,
109+ iArr , iArr + "8" , iArr + "16" , iArr + "32" , iArr + "64" ,
110+ fArr + "32" , fArr + "64" ,
111+ cArr + "64" , cArr + "128" ,
112+ uArr , uArr + "8" , uArr + "16" , uArr + "32" , uArr + "64" , uArr + "ptr" :
113+ arrayInit [jen .Id (n .Name )] = jen .Id (typeStr + "{}" )
114+ default :
115+ arrayInit [jen .Id (n .Name )] = jen .Id (string (src [pos .Type .Pos ()- 1 : pos .Type .End ()- 1 ])).Values (jen .Id (string (src [pos .Type .Pos ()+ 1 : pos .Type .End ()- 1 ])).Values ())
116+ }
84117 }
85118 }
86119 }
@@ -109,12 +142,12 @@ func GenGoCode(src []byte) (string, error) {
109142
110143 // Generate main function in new module.
111144 generatedCode .Func ().Id ("main" ).Params ().Block (init ... )
112- return fmt . Sprintf ( "%#v" , generatedCode ), nil
145+ return generatedCode . GoString ( ), nil
113146}
114147
115148// execGoCode executes and returns output from generated Go code.
116149func ExecGoCode (ctx context.Context , mainGo string ) ([]byte , error ) {
117- tmpDir , err := ioutil . TempDir ("" , "structgen" )
150+ tmpDir , err := os . MkdirTemp ("" , "structgen" )
118151 if err != nil {
119152 return nil , err
120153 }
0 commit comments