feat: add emit_str_enum config option

Author: devstein

internal/config.go

@@ -7,6 +7,7 @@ type Config struct {
 	Package                     string   `json:"package"`
 	Out                         string   `json:"out"`
 	EmitPydanticModels          bool     `json:"emit_pydantic_models"`
+	EmitStrEnum                 bool     `json:"emit_str_enum"`
 	QueryParameterLimit         *int32   `json:"query_parameter_limit"`
 	InflectionExcludeTableNames []string `json:"inflection_exclude_table_names"`
 }

internal/gen.go

@@ -681,12 +681,19 @@ func buildModelsTree(ctx *pyTmplCtx, i *importer) *pyast.Node {
 	mod.Body = append(mod.Body, buildImportGroup(std), buildImportGroup(pkg))
 
 	for _, e := range ctx.Enums {
+		bases := []*pyast.Node{
+			poet.Name("str"),
+			poet.Attribute(poet.Name("enum"), "Enum"),
+		}
+		if i.C.EmitStrEnum {
+			// override the bases to emit enum.StrEnum (only support in Python >=3.11)
+			bases = []*pyast.Node{
+				poet.Attribute(poet.Name("enum"), "StrEnum"),
+			}
+		}
 		def := &pyast.ClassDef{
-			Name: e.Name,
-			Bases: []*pyast.Node{
-				poet.Name("str"),
-				poet.Attribute(poet.Name("enum"), "Enum"),
-			},
+			Name:  e.Name,
+			Bases: bases,
 		}
 		if e.Comment != "" {
 			def.Body = append(def.Body, &pyast.Node{