support variadic types

coadler · coadler · commit a95194e60548 · 2022-05-10T14:19:36.000-05:00
diff --git a/coderd/audit/audit_gen.sh b/coderd/audit/audit_gen.sh
@@ -3,4 +3,4 @@
 set -euo pipefail
 
 cd "$(git rev-parse --show-toplevel)"
-go run ./scripts/auditgen ./coderd/database "$1"
+go run ./scripts/auditgen ./coderd/database "$@"
diff --git a/scripts/auditgen/main.go b/scripts/auditgen/main.go
@@ -19,7 +19,7 @@ import (
 func main() {
 	ctx := context.Background()
 	log := slog.Make(sloghuman.Sink(os.Stderr))
-	code, err := GenerateFromDirectory(ctx, os.Args[1], os.Args[2])
+	code, err := GenerateFromDirectory(ctx, os.Args[1], os.Args[2:]...)
 	if err != nil {
 		log.Fatal(ctx, "generate", slog.Error(err))
 	}
@@ -28,14 +28,14 @@ func main() {
 }
 
 // GenerateFromDirectory will return all the typescript code blocks for a directory
-func GenerateFromDirectory(ctx context.Context, directory, typName string) (string, error) {
+func GenerateFromDirectory(ctx context.Context, directory string, typeNames ...string) (string, error) {
 	g := Generator{}
 	err := g.parsePackage(ctx, directory)
 	if err != nil {
 		return "", xerrors.Errorf("parse package %q: %w", directory, err)
 	}
 
-	str, err := g.generate(typName)
+	str, err := g.generate(typeNames...)
 	if err != nil {
 		return "", xerrors.Errorf("parse package %q: %w", directory, err)
 	}
@@ -74,38 +74,40 @@ func (g *Generator) parsePackage(ctx context.Context, patterns ...string) error
 	return nil
 }
 
-func (g *Generator) generate(typName string) (string, error) {
+func (g *Generator) generate(typeNames ...string) (string, error) {
 	sb := strings.Builder{}
 
 	_, _ = fmt.Fprint(&sb, "Copy the following code into the audit.AuditableResources table\n\n")
 
-	obj := g.pkg.Types.Scope().Lookup(typName)
-	if obj == nil || obj.Type() == nil {
-		return "", xerrors.Errorf("type doesn't exist %q", typName)
-	}
-
-	switch obj := obj.(type) {
-	case *types.TypeName:
-		named, ok := obj.Type().(*types.Named)
-		if !ok {
-			panic("all typenames should be named types")
+	for _, typName := range typeNames {
+		obj := g.pkg.Types.Scope().Lookup(typName)
+		if obj == nil || obj.Type() == nil {
+			return "", xerrors.Errorf("type doesn't exist %q", typName)
 		}
 
-		switch typ := named.Underlying().(type) {
-		case *types.Struct:
-			g.writeStruct(&sb, typ, typName)
+		switch obj := obj.(type) {
+		case *types.TypeName:
+			named, ok := obj.Type().(*types.Named)
+			if !ok {
+				panic("all typenames should be named types")
+			}
+
+			switch typ := named.Underlying().(type) {
+			case *types.Struct:
+				g.writeStruct(&sb, typ, typName)
 
+			default:
+				return "", xerrors.Errorf("invalid type %T", obj)
+			}
 		default:
 			return "", xerrors.Errorf("invalid type %T", obj)
 		}
-	default:
-		return "", xerrors.Errorf("invalid type %T", obj)
 	}
 
 	return sb.String(), nil
 }
 
-func (g *Generator) writeStruct(w io.Writer, st *types.Struct, name string) {
+func (*Generator) writeStruct(w io.Writer, st *types.Struct, name string) {
 	_, _ = fmt.Fprintf(w, "\t&database.%s{}: {\n", name)
 
 	for i := 0; i < st.NumFields(); i++ {
