refactor: use poet.Node

QuentinN42 · QuentinN42 · commit 24fcd39f834e · 2025-07-25T18:07:43.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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/endtoend/testdata/emit_pydantic_models_with_fields/db/models.py b/internal/endtoend/testdata/emit_pydantic_models_with_fields/db/models.py
@@ -6,4 +6,6 @@
 
 class Author(pydantic.BaseModel):
     id: int
-    class: str
+    class_: str = pydantic.Field(
+        alias="class",
+    )
diff --git a/internal/endtoend/testdata/emit_pydantic_models_with_fields/db/query.py b/internal/endtoend/testdata/emit_pydantic_models_with_fields/db/query.py
@@ -26,7 +26,7 @@ def get_author(self, *, id: int) -> Optional[models.Author]:
             return None
         return models.Author(
             id=row[0],
-            class=row[1],
+            class_=row[1],
         )
 
 
@@ -40,5 +40,5 @@ async def get_author(self, *, id: int) -> Optional[models.Author]:
             return None
         return models.Author(
             id=row[0],
-            class=row[1],
+            class_=row[1],
         )
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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+      sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 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: "30dc4a9deda7f954e6444652e98763b3ad587ea672d88235b2c0f45ba8a4d03b"
+    sha256: "10f4fa000db23558eba8758926bab44a5e50118d3c892ecc4bd64a1fb88f678c"
 sql:
 - schema: schema.sql
   queries: query.sql
diff --git a/internal/gen.go b/internal/gen.go
@@ -107,7 +107,7 @@ func (v QueryValue) RowNode(rowVar string) *pyast.Node {
 	}
 	for i, f := range v.Struct.Fields {
 		call.Keywords = append(call.Keywords, &pyast.Keyword{
-			Arg: f.Name,
+			Arg: poet.FieldName(f.Name),
 			Value: subscriptNode(
 				rowVar,
 				constantInt(i),
@@ -622,51 +622,37 @@ func fieldNode(f Field, emitPydanticModels bool) *pyast.Node {
 		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,
-								},
-							},
+					Targets: poet.Nodes(
+						&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: poet.Node(
+						&pyast.Call{
+							Func: &pyast.Node{
+								Node: &pyast.Node_Attribute{
+									Attribute: &pyast.Attribute{
 										Value: &pyast.Node{
-											Node: &pyast.Node_Constant{
-												Constant: &pyast.Constant{
-													Value: &pyast.Constant_Str{
-														Str: f.Name,
-													},
+											Node: &pyast.Node_Name{
+												Name: &pyast.Name{
+													Id: "pydantic",
 												},
 											},
 										},
+										Attr: "Field",
 									},
 								},
 							},
+							Keywords: []*pyast.Keyword{
+								{
+									Arg:   "alias",
+									Value: poet.Constant(f.Name),
+								},
+							},
 						},
-					},
+					),
 				},
 			},
 		}

Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@ def get_author(self, *, id: int) -> Optional[models.Author]:`
`26`	`26`	`return None`
`27`	`27`	`return models.Author(`
`28`	`28`	`id=row[0],`
`29`		`- class=row[1],`
	`29`	`+ class_=row[1],`
`30`	`30`	`)`
`31`	`31`
`32`	`32`
`@@ -40,5 +40,5 @@ async def get_author(self, *, id: int) -> Optional[models.Author]:`
`40`	`40`	`return None`
`41`	`41`	`return models.Author(`
`42`	`42`	`id=row[0],`
`43`		`- class=row[1],`
	`43`	`+ class_=row[1],`
`44`	`44`	`)`