fix: Suport cross package types when Flattening (#427)

murarustefaan · web-flow · commit b79c286b8cbc · 2026-03-30T14:48:14.000Z
diff --git a/structs/flatten.go b/structs/flatten.go
@@ -46,6 +46,7 @@ type Field struct {
 type templateData struct {
 	HeaderComment string
 	PackageName   string
+	Imports       []string
 	SourceName    string
 	OutputName    string
 	Fields        []Field
@@ -100,6 +101,7 @@ func flattenOne(f *ast.File, cfg StructConfig, typeKinds map[string]string, scal
 	data := templateData{
 		HeaderComment: o.headerComment,
 		PackageName:   o.packageName,
+		Imports:       collectImports(fields),
 		SourceName:    cfg.SourceName,
 		OutputName:    cfg.OutputName,
 		Fields:        fields,
@@ -128,6 +130,38 @@ func flattenOne(f *ast.File, cfg StructConfig, typeKinds map[string]string, scal
 	return nil
 }
 
+// collectImports extracts unique package import paths from field types that
+// contain a dot (e.g. "time.Time" -> "time"). Returns a sorted, deduplicated list.
+func collectImports(fields []Field) []string {
+	// Map well-known qualified type prefixes to their import paths.
+	pkgToImport := map[string]string{
+		"time": "time",
+		"net":  "net",
+		"url":  "net/url",
+		"json": "encoding/json",
+		"uuid": "github.com/google/uuid",
+	}
+
+	seen := map[string]bool{}
+	for _, f := range fields {
+		typ := strings.TrimPrefix(f.Type, "*")
+		typ = strings.TrimPrefix(typ, "[]")
+		if dot := strings.IndexByte(typ, '.'); dot > 0 {
+			pkg := typ[:dot]
+			if imp, ok := pkgToImport[pkg]; ok && !seen[imp] {
+				seen[imp] = true
+			}
+		}
+	}
+
+	imports := make([]string, 0, len(seen))
+	for imp := range seen {
+		imports = append(imports, imp)
+	}
+	slices.Sort(imports)
+	return imports
+}
+
 // sortFields sorts fields with "ID" first, then alphabetically by name (case-insensitive).
 func sortFields(fields []Field) {
 	slices.SortStableFunc(fields, func(a, b Field) int {
diff --git a/structs/flatten_test.go b/structs/flatten_test.go
@@ -202,8 +202,13 @@ func TestResolveType(t *testing.T) {
 			want: "[]map[string]any",
 		},
 		{
-			name: "cross-package type",
+			name: "cross-package scalar (time.Time)",
 			expr: &ast.SelectorExpr{X: &ast.Ident{Name: "time"}, Sel: &ast.Ident{Name: "Time"}},
+			want: "time.Time",
+		},
+		{
+			name: "cross-package struct",
+			expr: &ast.SelectorExpr{X: &ast.Ident{Name: "api"}, Sel: &ast.Ident{Name: "Finding"}},
 			want: "map[string]any",
 		},
 		{
diff --git a/structs/options.go b/structs/options.go
@@ -1,12 +1,35 @@
 package structs
 
 // defaultScalarKinds are Go types that are inexpensive to decode and should stay typed.
+// Includes both unqualified names (for types in the same package) and
+// fully qualified names like "time.Time" (for cross-package types).
 var defaultScalarKinds = map[string]bool{
 	"string":  true,
 	"bool":    true,
+	"float32": true,
 	"float64": true,
-	"int64":   true,
 	"int":     true,
+	"int8":    true,
+	"int16":   true,
+	"int32":   true,
+	"int64":   true,
+	"uint":    true,
+	"uint8":   true,
+	"uint16":  true,
+	"uint32":  true,
+	"uint64":  true,
+	"byte":    true,
+	"rune":    true,
+
+	"time.Time":        true,
+	"time.Duration":    true,
+	"net.IP":           true,
+	"net.HardwareAddr": true,
+	"url.URL":          true,
+	"json.RawMessage":  true,
+	"json.Number":      true,
+
+	"uuid.UUID": true,
 }
 
 type options struct {
diff --git a/structs/resolve.go b/structs/resolve.go
@@ -80,8 +80,18 @@ func resolveType(expr ast.Expr, typeKinds map[string]string, scalars map[string]
 		}
 		return "[]" + inner
 
+	case *ast.SelectorExpr:
+		// Cross-package types like time.Time.
+		if pkg, ok := t.X.(*ast.Ident); ok {
+			qualified := pkg.Name + "." + t.Sel.Name
+			if scalars[qualified] {
+				return qualified
+			}
+		}
+		return "map[string]any"
+
 	default:
-		// SelectorExpr (cross-package types), MapType, InterfaceType, etc.
+		// MapType, InterfaceType, etc.
 		return "map[string]any"
 	}
 }
diff --git a/structs/templates/struct.go.tpl b/structs/templates/struct.go.tpl
@@ -1,7 +1,13 @@
 // {{ .HeaderComment }}
 
 package {{ .PackageName }}
-
+{{ if .Imports }}
+import (
+{{- range .Imports }}
+	"{{ . }}"
+{{- end }}
+)
+{{ end }}
 // {{ .OutputName }} is an optimized version of {{ .SourceName }}
 // where deeply nested struct fields are replaced with map[string]any to avoid
 // the decode -> allocate -> re-encode cycle.

Original file line number	Diff line number	Diff line change
`@@ -80,8 +80,18 @@ func resolveType(expr ast.Expr, typeKinds map[string]string, scalars map[string]`
`80`	`80`	`}`
`81`	`81`	`return "[]" + inner`
`82`	`82`
	`83`	`+ case *ast.SelectorExpr:`
	`84`	`+ // Cross-package types like time.Time.`
	`85`	`+ if pkg, ok := t.X.(*ast.Ident); ok {`
	`86`	`+ qualified := pkg.Name + "." + t.Sel.Name`
	`87`	`+ if scalars[qualified] {`
	`88`	`+ return qualified`
	`89`	`+ }`
	`90`	`+ }`
	`91`	`+ return "map[string]any"`
	`92`	`+`
`83`	`93`	`default:`
`84`		`- // SelectorExpr (cross-package types), MapType, InterfaceType, etc.`
	`94`	`+ // MapType, InterfaceType, etc.`
`85`	`95`	`return "map[string]any"`
`86`	`96`	`}`
`87`	`97`	`}`