fix: added reserved python keywords

QuentinN42 · QuentinN42 · commit f92824439ae1 · 2025-07-25T17:55:11.000+02:00
Signed-off-by: QuentinN42 &lt;quentin@lieumont.fr&gt;
diff --git a/internal/endtoend/testdata/emit_pydantic_models/sqlc.yaml b/internal/endtoend/testdata/emit_pydantic_models/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/emit_pydantic_models_with_fields/sqlc.yaml b/internal/endtoend/testdata/emit_pydantic_models_with_fields/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/emit_str_enum/sqlc.yaml b/internal/endtoend/testdata/emit_str_enum/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
   - name: py
     wasm:
       url: file://../../../../bin/sqlc-gen-python.wasm
-      sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+      sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
   - schema: schema.sql
     queries: query.sql
diff --git a/internal/endtoend/testdata/exec_result/sqlc.yaml b/internal/endtoend/testdata/exec_result/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/exec_rows/sqlc.yaml b/internal/endtoend/testdata/exec_rows/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/inflection_exclude_table_names/sqlc.yaml b/internal/endtoend/testdata/inflection_exclude_table_names/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/query_parameter_limit_two/sqlc.yaml b/internal/endtoend/testdata/query_parameter_limit_two/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/query_parameter_limit_undefined/sqlc.yaml b/internal/endtoend/testdata/query_parameter_limit_undefined/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/query_parameter_limit_zero/sqlc.yaml b/internal/endtoend/testdata/query_parameter_limit_zero/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/query_parameter_no_limit/sqlc.yaml b/internal/endtoend/testdata/query_parameter_no_limit/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "b220b970c2b429b3d3cdde49d2f49c5db73d85c98253ec053533822ab2b58640"
+    sha256: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/gen.go b/internal/gen.go
@@ -601,11 +601,82 @@ func pydanticNode(name string) *pyast.ClassDef {
 	}
 }
 
-func fieldNode(f Field) *pyast.Node {
+func fieldNode(f Field, emitPydanticModels bool) *pyast.Node {
+	if !poet.IsReserved(f.Name) {
+		return &pyast.Node{
+			Node: &pyast.Node_AnnAssign{
+				AnnAssign: &pyast.AnnAssign{
+					Target:     &pyast.Name{Id: f.Name},
+					Annotation: f.Type.Annotation(),
+					Comment:    f.Comment,
+				},
+			},
+		}
+	}
+
+	// At this point the field name is a reserved python keyword, so we need to
+	// update the field name to be a valid python identifier.
+	if emitPydanticModels {
+		// On Pydantic add `_` at the end of the field name to make it valid
+		// Also add a `= pydantic.Field(alias=...)` to the field to allow clean serde
+		return &pyast.Node{
+			Node: &pyast.Node_Assign{
+				Assign: &pyast.Assign{
+					Targets: []*pyast.Node{
+						{
+							Node: &pyast.Node_AnnAssign{
+								AnnAssign: &pyast.AnnAssign{
+									Target:     &pyast.Name{Id: poet.FieldName(f.Name)},
+									Annotation: f.Type.Annotation(),
+									Comment:    f.Comment,
+								},
+							},
+						},
+					},
+					Value: &pyast.Node{
+						Node: &pyast.Node_Call{
+							Call: &pyast.Call{
+								Func: &pyast.Node{
+									Node: &pyast.Node_Attribute{
+										Attribute: &pyast.Attribute{
+											Value: &pyast.Node{
+												Node: &pyast.Node_Name{
+													Name: &pyast.Name{
+														Id: "pydantic",
+													},
+												},
+											},
+											Attr: "Field",
+										},
+									},
+								},
+								Keywords: []*pyast.Keyword{
+									{
+										Arg: "alias",
+										Value: &pyast.Node{
+											Node: &pyast.Node_Constant{
+												Constant: &pyast.Constant{
+													Value: &pyast.Constant_Str{
+														Str: f.Name,
+													},
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+	}
+
+	// On dataclasses add `_` at the end of the field name to make it valid
 	return &pyast.Node{
 		Node: &pyast.Node_AnnAssign{
 			AnnAssign: &pyast.AnnAssign{
-				Target:     &pyast.Name{Id: f.Name},
+				Target:     &pyast.Name{Id: poet.FieldName(f.Name)},
 				Annotation: f.Type.Annotation(),
 				Comment:    f.Comment,
 			},
@@ -731,7 +802,7 @@ func buildModelsTree(ctx *pyTmplCtx, i *importer) *pyast.Node {
 			})
 		}
 		for _, f := range m.Fields {
-			def.Body = append(def.Body, fieldNode(f))
+			def.Body = append(def.Body, fieldNode(f, ctx.C.EmitPydanticModels))
 		}
 		mod.Body = append(mod.Body, &pyast.Node{
 			Node: &pyast.Node_ClassDef{
@@ -857,7 +928,7 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 					def = dataclassNode(arg.Struct.Name)
 				}
 				for _, f := range arg.Struct.Fields {
-					def.Body = append(def.Body, fieldNode(f))
+					def.Body = append(def.Body, fieldNode(f, ctx.C.EmitPydanticModels))
 				}
 				mod.Body = append(mod.Body, poet.Node(def))
 			}
@@ -870,7 +941,7 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 				def = dataclassNode(q.Ret.Struct.Name)
 			}
 			for _, f := range q.Ret.Struct.Fields {
-				def.Body = append(def.Body, fieldNode(f))
+				def.Body = append(def.Body, fieldNode(f, ctx.C.EmitPydanticModels))
 			}
 			mod.Body = append(mod.Body, poet.Node(def))
 		}
diff --git a/internal/gen_test.go b/internal/gen_test.go
@@ -0,0 +1,53 @@
+package python
+
+import (
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/sqlc-dev/sqlc-gen-python/internal/printer"
+)
+
+func Test_fieldNode(t *testing.T) {
+	cases := []struct {
+		Expected           string
+		EmitPydanticModels bool
+		Field              Field
+	}{
+		{
+			Expected: "class_: int",
+			Field: Field{
+				Name: "class",
+				Type: pyType{
+					InnerType: "int",
+					IsNull:    false,
+					IsArray:   false,
+				},
+			},
+			EmitPydanticModels: false,
+		},
+		{
+			Expected: "class_: int = pydantic.Field(\n    alias=\"class\",\n)",
+			Field: Field{
+				Name: "class",
+				Type: pyType{
+					InnerType: "int",
+					IsNull:    false,
+					IsArray:   false,
+				},
+			},
+			EmitPydanticModels: true,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.Expected+" "+strconv.FormatBool(tc.EmitPydanticModels), func(t *testing.T) {
+			res := fieldNode(tc.Field, tc.EmitPydanticModels)
+			result := printer.Print(res, printer.Options{})
+			if diff := cmp.Diff(strings.TrimSpace(tc.Expected), strings.TrimSpace(string(result.Python))); diff != "" {
+				t.Errorf("node to python code mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
diff --git a/internal/poet/builders.go b/internal/poet/builders.go
@@ -64,7 +64,7 @@ func Is() *ast.Node {
 func Name(id string) *ast.Node {
 	return &ast.Node{
 		Node: &ast.Node_Name{
-			Name: &ast.Name{Id: id},
+			Name: &ast.Name{Id: FieldName(id)},
 		},
 	}
 }
diff --git a/internal/poet/reserved.go b/internal/poet/reserved.go
@@ -0,0 +1,38 @@
+package poet
+
+import "slices"
+
+// TODO(quentin@escape.tech): check if this is complete
+var reservedKeywords = []string{
+	"class",
+	"if",
+	"else",
+	"elif",
+	"not",
+	"for",
+	"and",
+	"in",
+	"is",
+	"or",
+	"with",
+	"as",
+	"assert",
+	"break",
+	"except",
+	"finally",
+	"try",
+	"raise",
+	"return",
+	"yield",
+}
+
+func IsReserved(name string) bool {
+	return slices.Contains(reservedKeywords, name)
+}
+
+func FieldName(name string) string {
+	if IsReserved(name) {
+		return name + "_"
+	}
+	return name
+}
