mkunion serde, go:tag support, simplification of x/shape

cmd/mkfunc/main.go

@@ -41,7 +41,7 @@ func main() {
 
 			match := generators.FunctionMatchGenerator{
 				Header:      "// Code generated by mkfunc. DO NOT EDIT.",
-				PackageName: inferred.PackageName,
+				PackageName: inferred.PackageName(),
 				MaxSize:     c.Int("max-size"),
 			}
 

cmd/mkunion/main.go

@@ -16,6 +16,16 @@ import (
 	"syscall"
 )
 
+type extractImports interface {
+	ExtractImports(x shape.Shape) generators.PkgMap
+	SkipImportsAndPackage(x bool)
+}
+
+type extractInitFuncs interface {
+	ExtractImportFuncs(s shape.Shape) []string
+	SkipInitFunc(flag bool)
+}
+
 func main() {
 	ctx, _ := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 
@@ -93,42 +103,38 @@ func main() {
 					return err
 				}
 
-				unionNames := c.StringSlice("name")
-				if len(unionNames) == 0 {
-					unionNames = inferred.PossibleUnionTypes()
-				}
-
-				for _, unionName := range unionNames {
-					options := []generators.GenerateOption{
-						generators.WithPackageName(inferred.PackageName),
+				var unions []*shape.UnionLike
+				for _, unionName := range c.StringSlice("name") {
+					union := inferred.RetrieveUnion(unionName)
+					if union == nil {
+						return fmt.Errorf("union %s not found in %s", unionName, sourcePath)
 					}
 
-					if !c.Bool("no-compact") {
-						options = append(options, generators.WithBufferedImports())
-					}
+					unions = append(unions, union)
+				}
 
-					helper := generators.NewHelper(options...)
-					union := inferred.RetrieveUnion(unionName)
+				if len(unions) == 0 {
+					unions = inferred.RetrieveUnions()
+				}
+
+				for _, union := range unions {
 					if union == nil {
-						return fmt.Errorf("union %s not found in %s", unionName, sourcePath)
+						return fmt.Errorf("union %s not found in %s", union.Name, sourcePath)
 					}
 
-					jsonGenerator := generators.NewDeSerJSONGenerator(union, helper)
-					shapeGenerator := generators.NewShapeGenerator(union, helper)
-					visitor := generators.NewVisitorGenerator(union, helper)
-					schema := generators.NewSchemaGenerator(union, helper)
+					jsonGenerator := generators.SerdeJSONUnion(union)
+					shapeGenerator := generators.NewShapeUnion(union)
+					visitor := generators.NewVisitorGenerator(union)
 
-					// ensures that order of generators is always the same
+					// ensures that order of generators2 is always the same
 					generatorsList := []string{
 						"visitor",
-						"schema",
 						"shape",
 						"json",
 					}
 
-					generators := map[string]generators.Generator{
+					generators2 := map[string]generators.Generator{
 						"visitor": visitor,
-						"schema":  schema,
 						"shape":   shapeGenerator,
 						"json":    jsonGenerator,
 					}
@@ -148,12 +154,12 @@ func main() {
 
 					for _, name := range skipExtension {
 						log.Infof("skip extension %s", name)
-						delete(generators, name)
+						delete(generators2, name)
 					}
 
 					if c.Bool("no-compact") {
 						for _, name := range generatorsList {
-							g, ok := generators[name]
+							g, ok := generators2[name]
 							if !ok {
 								continue
 							}
@@ -172,35 +178,58 @@ func main() {
 							}
 						}
 					} else {
-						body := bytes.Buffer{}
+						packageName := union.PkgName
+						pkgMap := make(generators.PkgMap)
+						initFunc := make(generators.InitFuncs, 0, 0)
+						shapesContents := bytes.Buffer{}
+
 						for _, name := range generatorsList {
-							g, ok := generators[name]
+							g, ok := generators2[name]
 							if !ok {
 								continue
 							}
+							if gen, ok := g.(extractImports); ok {
+								gen.SkipImportsAndPackage(true)
+							}
+							if gen, ok := g.(extractInitFuncs); ok {
+								gen.SkipInitFunc(true)
+							}
 
 							b, err := g.Generate()
 							if err != nil {
 								return fmt.Errorf("failed to generate %s for %s in %s: %w", name, union.Name, sourcePath, err)
 							}
-							body.WriteString(fmt.Sprintf("//mkunion-extension:%s\n", name))
-							body.Write(b)
-							body.WriteString("\n")
-						}
 
-						header := bytes.Buffer{}
-						header.WriteString(helper.RenderBufferedHeader())
-						header.WriteString(helper.RenderBufferedImport())
-						log.Infof(helper.RenderBufferedImport())
+							if gen, ok := g.(extractImports); ok {
+								gen.ExtractImports(union)
+								pkgMap = generators.MergePkgMaps(pkgMap,
+									gen.ExtractImports(union),
+								)
+							}
+							if gen, ok := g.(extractInitFuncs); ok {
+								initFunc = append(initFunc, gen.ExtractImportFuncs(union)...)
+							}
 
-						fileName := baseName + "_" + strings.ToLower(union.Name) + "_gen.go"
-						log.Infof("writing %s", fileName)
+							shapesContents.WriteString(fmt.Sprintf("//mkunion-extension:%s\n", name))
+							shapesContents.Write(b)
+							shapesContents.WriteString("\n")
+						}
 
-						header.Write(body.Bytes())
+						contents := "// Code generated by mkunion. DO NOT EDIT.\n"
+						contents += fmt.Sprintf("package %s\n\n", packageName)
+						contents += generators.GenerateImports(pkgMap)
+						contents += generators.GenerateInitFunc(initFunc)
+						contents += shapesContents.String()
 
-						err = os.WriteFile(path.Join(cwd, fileName), header.Bytes(), 0644)
+						fileName := path.Join(
+							path.Dir(sourcePath),
+							fmt.Sprintf("%s_%s_gen.go", baseName, strings.ToLower(union.Name)),
+						)
+
+						log.Infof("writing %s", fileName)
+						err = os.WriteFile(fileName, []byte(contents), 0644)
 						if err != nil {
-							return fmt.Errorf("failed to write %s for %s in %s: %w", "gen", union.Name, sourcePath, err)
+							return fmt.Errorf("failed to write(2) %s for %s in %s: %w", "gen", union.Name, sourcePath, err)
 						}
 					}
 				}
@@ -257,10 +286,123 @@ func main() {
 				},
 			},
 			{
-				Name: "shape-export",
+				Name: "serde",
 				Flags: []cli.Flag{
 					&cli.StringFlag{
 						Name:        "type",
+						DefaultText: "json",
+					},
+					&cli.StringSliceFlag{
+						Name:      "input-go-file",
+						Aliases:   []string{"i", "input"},
+						Usage:     `When not provided, it will try to use GOFILE environment variable, used when combined with //go:generate mkunion -name=MyUnionType`,
+						TakesFile: true,
+					},
+					&cli.BoolFlag{
+						Name:     "verbose",
+						Aliases:  []string{"v"},
+						Required: false,
+						Value:    false,
+					},
+				},
+				Action: func(c *cli.Context) error {
+					if c.Bool("verbose") {
+						log.SetLevel(log.DebugLevel)
+					}
+
+					sourcePaths := c.StringSlice("input-go-file")
+					if len(sourcePaths) == 0 && os.Getenv("GOFILE") != "" {
+						cwd, _ := syscall.Getwd()
+						sourceName := path.Base(os.Getenv("GOFILE"))
+						sourcePaths = []string{
+							path.Join(cwd, sourceName),
+						}
+					}
+
+					if len(sourcePaths) == 0 {
+						// show usage
+						cli.ShowAppHelpAndExit(c, 1)
+					}
+
+					for _, sourcePath := range sourcePaths {
+						inferred, err := shape.InferFromFile(sourcePath)
+						if err != nil {
+							return fmt.Errorf("failed inferring shape in %s; %w", sourcePath, err)
+						}
+
+						shapes := inferred.RetrieveShapesTaggedAs("serde")
+						if len(shapes) == 0 {
+							log.Infof("no shape found in %s", sourcePath)
+							continue
+						}
+
+						packageName := "main"
+						pkgMap := make(generators.PkgMap)
+						initFunc := make(generators.InitFuncs, 0, 0)
+						shapesContents := bytes.Buffer{}
+
+						for _, x := range shapes {
+							packageName = shape.ToGoPkgName(x)
+							genSerde := generators.NewSerdeJSONTagged(x)
+							genSerde.SkipImportsAndPackage(true)
+
+							genShape := generators.NewShapeTagged(x)
+							genShape.SkipImportsAndPackage(true)
+							genShape.SkipInitFunc(true)
+
+							contents := "//serde:json\n"
+							contents, err = genSerde.Generate()
+							if err != nil {
+								return fmt.Errorf("failed to generate json serde for %s in %s: %w", shape.ToGoTypeName(x), sourcePath, err)
+							}
+							shapesContents.WriteString(contents)
+
+							contents = "//shape\n"
+							contents, err = genShape.Generate()
+							if err != nil {
+								return fmt.Errorf("failed to generate shape for %s in %s: %w", shape.ToGoTypeName(x), sourcePath, err)
+							}
+							shapesContents.WriteString(contents)
+
+							pkgMap = generators.MergePkgMaps(pkgMap,
+								genSerde.ExtractImports(x),
+								genShape.ExtractImports(x),
+							)
+
+							initFunc = append(initFunc, genShape.ExtractImportFuncs(x)...)
+						}
+
+						contents := "// Code generated by mkunion. DO NOT EDIT.\n"
+						contents += fmt.Sprintf("package %s\n\n", packageName)
+						contents += generators.GenerateImports(pkgMap)
+						contents += generators.GenerateInitFunc(initFunc)
+						contents += shapesContents.String()
+
+						sourceName := path.Base(sourcePath)
+						baseName := strings.TrimSuffix(sourceName, path.Ext(sourceName))
+						fileName := path.Join(
+							path.Dir(sourcePath),
+							fmt.Sprintf("%s_serde_gen.go", baseName),
+						)
+
+						log.Infof("writing %s", fileName)
+						err = os.WriteFile(fileName, []byte(contents), 0644)
+						if err != nil {
+							return fmt.Errorf("failed to write serialiser in %s: %w", sourcePath, err)
+						}
+
+						return nil
+					}
+
+					return nil
+				},
+			},
+			{
+				Name: "shape-export",
+				Flags: []cli.Flag{
+					&cli.StringFlag{
+						Name:        "language",
+						Aliases:     []string{"lang"},
 						DefaultText: "typescript",
 					},
 					&cli.StringFlag{
@@ -273,8 +415,18 @@ func main() {
 						Usage:     `When not provided, it will try to use GOFILE environment variable, used when combined with //go:generate mkunion -name=MyUnionType`,
 						TakesFile: true,
 					},
+					&cli.BoolFlag{
+						Name:     "verbose",
+						Aliases:  []string{"v"},
+						Required: false,
+						Value:    false,
+					},
 				},
 				Action: func(c *cli.Context) error {
+					if c.Bool("verbose") {
+						log.SetLevel(log.DebugLevel)
+					}
+
 					sourcePaths := c.StringSlice("input-go-file")
 					if len(sourcePaths) == 0 && os.Getenv("GOFILE") != "" {
 						cwd, _ := syscall.Getwd()
@@ -297,15 +449,14 @@ func main() {
 							return err
 						}
 
-						for _, union := range inferred.RetrieveUnions() {
-							tsr.AddUnion(union)
-						}
-
-						for _, structLike := range inferred.RetrieveStructs() {
-							tsr.AddStruct(structLike)
+						for _, x := range inferred.RetrieveShapes() {
+							tsr.AddShape(x)
+							tsr.FollowRef(x)
 						}
 					}
 
+					tsr.FollowImports()
+
 					err := tsr.WriteToDir(c.String("output-dir"))
 					if err != nil {
 						return fmt.Errorf("failed to write to dir %s: %w", c.String("output-dir"), err)

example/my-app/functions.go

@@ -12,6 +12,7 @@ import (
 	"image/jpeg"
 	"image/png"
 	"io"
+	"math/rand"
 	"os"
 )
 
@@ -31,6 +32,25 @@ var functions = map[string]workflow.Function{
 			Result: schema.MkString(a + b),
 		}, nil
 	},
+	"concat_error": func(body *workflow.FunctionInput) (*workflow.FunctionOutput, error) {
+		args := body.Args
+		a, ok := schema.As[string](args[0])
+		if !ok {
+			return nil, fmt.Errorf("expected string, got %T", args[0])
+		}
+		b, ok := schema.As[string](args[1])
+		if !ok {
+			return nil, fmt.Errorf("expected string, got %T", args[1])
+		}
+
+		if rand.Float32() > 0.5 {
+			return nil, fmt.Errorf("random error")
+		}
+
+		return &workflow.FunctionOutput{
+			Result: schema.MkString(a + b),
+		}, nil
+	},
 	"genimageb64": func(body *workflow.FunctionInput) (*workflow.FunctionOutput, error) {
 		args := body.Args
 		_, ok := schema.As[string](args[0])