Preserve original value type

Author: codemageddon

_examples/status/job_status_enum.go

@@ -13,7 +13,7 @@ import (
 // {{.Type | title}} is the exported type for the enum
 type {{.Type | title}} struct {
 	name  string
-	value int
+	value {{.OriginalType}}
 }
 
 func (e {{.Type | title}}) String() string { return e.name }

_examples/status/status_enum.go

@@ -26,11 +26,12 @@ var titleCaser = cases.Title(language.English, cases.NoLower)
 
 // Generator holds the data needed for enum code generation
 type Generator struct {
-	Type      string         // the private type name (e.g., "status")
-	Path      string         // output directory path
-	values    map[string]int // const values found
-	pkgName   string         // package name from source file
-	lowerCase bool           // use lower case for marshal/unmarshal
+	Type         string         // the private type name (e.g., "status")
+	Path         string         // output directory path
+	values       map[string]int // const values found
+	pkgName      string         // package name from source file
+	lowerCase    bool           // use lower case for marshal/unmarshal
+	originalType string         // original type name (e.g., "uint8")
 }
 
 // Value represents a single enum value
@@ -81,6 +82,9 @@ func (g *Generator) Parse(dir string) error {
 		}
 	}
 
+	if g.originalType == "" {
+		return fmt.Errorf("type %s not found", g.Type)
+	}
 	if len(g.values) == 0 {
 		return fmt.Errorf("no const values found for type %s", g.Type)
 	}
@@ -90,7 +94,22 @@ func (g *Generator) Parse(dir string) error {
 
 // parseFile processes a single file for enum declarations
 func (g *Generator) parseFile(file *ast.File) {
-
+	parseTypeBlock := func(decl *ast.GenDecl) {
+		// extracts the type name from a const block
+		for _, spec := range decl.Specs {
+			vspec, ok := spec.(*ast.TypeSpec)
+			if !ok {
+				continue
+			}
+			if vspec.Name.Name == g.Type {
+				tspec, ok := vspec.Type.(*ast.Ident)
+				if !ok {
+					continue
+				}
+				g.originalType = tspec.Name
+			}
+		}
+	}
 	parseConstBlock := func(decl *ast.GenDecl) {
 		// extracts enum values from a const block
 		var iotaVal int
@@ -148,8 +167,13 @@ func (g *Generator) parseFile(file *ast.File) {
 	}
 
 	ast.Inspect(file, func(n ast.Node) bool {
-		if decl, ok := n.(*ast.GenDecl); ok && decl.Tok == token.CONST {
-			parseConstBlock(decl)
+		if decl, ok := n.(*ast.GenDecl); ok {
+			switch decl.Tok {
+			case token.CONST:
+				parseConstBlock(decl)
+			case token.TYPE:
+				parseTypeBlock(decl)
+			}
 		}
 		return true
 	})
@@ -216,15 +240,17 @@ func (g *Generator) Generate() error {
 
 	// prepare template data
 	data := struct {
-		Type      string
-		Values    []Value
-		Package   string
-		LowerCase bool
+		Type         string
+		Values       []Value
+		Package      string
+		LowerCase    bool
+		OriginalType string
 	}{
-		Type:      g.Type,
-		Values:    values,
-		Package:   pkgName,
-		LowerCase: g.lowerCase,
+		Type:         g.Type,
+		Values:       values,
+		Package:      pkgName,
+		LowerCase:    g.lowerCase,
+		OriginalType: g.originalType,
 	}
 
 	// execute template

internal/generator/enum.go.tmpl

@@ -368,6 +368,7 @@ func TestPermissions(t *testing.T) {
 
 		// create a sample status file
 		sampleFile := `package source
+type status uint8
 const (
 	statusUnknown = iota
 	statusActive
@@ -466,6 +467,7 @@ func TestParseSpecialCases(t *testing.T) {
 		tmpDir := t.TempDir()
 		err := os.WriteFile(filepath.Join(tmpDir, "empty.go"), []byte(`
 package test
+type status uint8
 const (
 )
 `), 0o644)
@@ -483,6 +485,7 @@ const (
 		tmpDir := t.TempDir()
 		err := os.WriteFile(filepath.Join(tmpDir, "no_values.go"), []byte(`
 package test
+type status uint8
 const name string
 `), 0o644)
 		require.NoError(t, err)
@@ -494,6 +497,26 @@ const name string
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "no const values found for type status")
 	})
+
+	t.Run("no status type", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		err := os.WriteFile(filepath.Join(tmpDir, "no_type.go"), []byte(`
+package test
+const (
+	statusUnknown = iota
+	statusActive
+	statusInactive
+)
+`), 0o644)
+		require.NoError(t, err)
+
+		gen, err := New("status", "")
+		require.NoError(t, err)
+
+		err = gen.Parse(tmpDir)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "type status not found")
+	})
 }
 
 func TestSplitCamelCase(t *testing.T) {