feat: with Mapped class annotation

zztkm · zztkm · commit 48821ef7a437 · 2023-09-06T14:46:38.000+09:00
diff --git a/_examples/gen/sqlc/orm.py b/_examples/gen/sqlc/orm.py
@@ -1,17 +1,18 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.20.0
-import sqlalchemy
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 from typing import Optional
 
 
-class Base(sqlalchemy.orm.DeclarativeBase):
+class Base(DeclarativeBase):
     pass
 
 
-class Author(hello.Base):
-    id: int
-    name: str
-    age: int
-    bio: Optional[str]
-    is_active: bool
+class Author(Base):
+    __tablename__ = "authors"
+    id: Mapped[int]
+    name: Mapped[str]
+    age: Mapped[int]
+    bio: Mapped[Optional[str]]
+    is_active: Mapped[bool]
diff --git a/internal/gen.go b/internal/gen.go
@@ -623,15 +623,8 @@ func sqlalchemyNode(name string) *pyast.ClassDef {
 		Name: name,
 		Bases: []*pyast.Node{
 			{
-				Node: &pyast.Node_Attribute{
-					Attribute: &pyast.Attribute{
-						Value: &pyast.Node{
-							Node: &pyast.Node_Name{
-								Name: &pyast.Name{Id: "hello"},
-							},
-						},
-						Attr: "Base",
-					},
+				Node: &pyast.Node_Name{
+					Name: &pyast.Name{Id: "Base"},
 				},
 			},
 		},
@@ -643,15 +636,8 @@ func sqlalchemyBaseNode() *pyast.ClassDef {
 		Name: "Base",
 		Bases: []*pyast.Node{
 			{
-				Node: &pyast.Node_Attribute{
-					Attribute: &pyast.Attribute{
-						Value: &pyast.Node{
-							Node: &pyast.Node_Name{
-								Name: &pyast.Name{Id: "sqlalchemy.orm"},
-							},
-						},
-						Attr: "DeclarativeBase",
-					},
+				Node: &pyast.Node_Name{
+					Name: &pyast.Name{Id: "DeclarativeBase"},
 				},
 			},
 		},
@@ -670,6 +656,18 @@ func fieldNode(f Field) *pyast.Node {
 	}
 }
 
+func ormFieldNode(f Field) *pyast.Node {
+	return &pyast.Node{
+		Node: &pyast.Node_AnnAssign{
+			AnnAssign: &pyast.AnnAssign{
+				Target:     &pyast.Name{Id: f.Name},
+				Annotation: subscriptNode("Mapped", f.Type.Annotation()),
+				Comment:    f.Comment,
+			},
+		},
+	}
+}
+
 func fieldPassNode() *pyast.Node {
 	return &pyast.Node{
 		Node: &pyast.Node_Pass{
@@ -716,6 +714,7 @@ func buildImportGroup(specs map[string]importSpec) *pyast.Node {
 	var body []*pyast.Node
 	for _, spec := range buildImportBlock2(specs) {
 		if len(spec.Names) > 0 && spec.Names[0] != "" {
+			// e.g. from sqlalchemy import create_engine
 			imp := &pyast.ImportFrom{
 				Module: spec.Module,
 			}
@@ -834,6 +833,7 @@ func buildOrmTree(ctx *pyTmplCtx, i *importer) *pyast.Node {
 	}
 
 	// generate Base class
+	// Base class is a dependency for all orm models
 	def := sqlalchemyBaseNode()
 	def.Body = append(def.Body, fieldPassNode())
 	mod.Body = append(mod.Body, &pyast.Node{
@@ -842,7 +842,7 @@ func buildOrmTree(ctx *pyTmplCtx, i *importer) *pyast.Node {
 		},
 	})
 
-	for _, m := range ctx.Orms {
+	for _, m := range ctx.Models {
 		def := sqlalchemyNode(m.Name)
 		if m.Comment != "" {
 			def.Body = append(def.Body, &pyast.Node{
@@ -853,8 +853,11 @@ func buildOrmTree(ctx *pyTmplCtx, i *importer) *pyast.Node {
 				},
 			})
 		}
+		// add table name as class attribute
+		def.Body = append(def.Body, assignNode("__tablename__", poet.Constant(m.Table.Name)))
+
 		for _, f := range m.Fields {
-			def.Body = append(def.Body, fieldNode(f))
+			def.Body = append(def.Body, ormFieldNode(f))
 		}
 		mod.Body = append(mod.Body, &pyast.Node{
 			Node: &pyast.Node_ClassDef{
@@ -1190,7 +1193,6 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 type pyTmplCtx struct {
 	SqlcVersion string
 	Models      []Struct
-	Orms        []Struct
 	Queries     []Query
 	Enums       []Enum
 	SourceName  string
@@ -1230,7 +1232,6 @@ func Generate(_ context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenRes
 
 	tctx := pyTmplCtx{
 		Models:      models,
-		Orms:        models,
 		Queries:     queries,
 		Enums:       enums,
 		SqlcVersion: req.SqlcVersion,
diff --git a/internal/imports.go b/internal/imports.go
@@ -121,6 +121,7 @@ func (i *importer) modelImportSpecs() (map[string]importSpec, map[string]importS
 	return std, pkg
 }
 
+// ormImportSpecs returns the standard and package imports for the ORM.
 func (i *importer) ormImportSpecs() (map[string]importSpec, map[string]importSpec) {
 	modelUses := func(name string) bool {
 		for _, model := range i.Models {
@@ -133,7 +134,9 @@ func (i *importer) ormImportSpecs() (map[string]importSpec, map[string]importSpe
 
 	std := stdImports(modelUses)
 	if i.C.EmitSQLAlchemyModels {
-		std["sqlalchemy"] = importSpec{Module: "sqlalchemy"}
+		std["sqlalchemy.orm.DeclarativeBase"] = importSpec{Module: "sqlalchemy.orm", Name: "DeclarativeBase"}
+		std["sqlalchemy.orm.Mapped"] = importSpec{Module: "sqlalchemy.orm", Name: "Mapped"}
+		std["sqlalchemy.orm.mapped_column"] = importSpec{Module: "sqlalchemy.orm", Name: "mapped_column"}
 	}
 	if len(i.Enums) > 0 {
 		std["enum"] = importSpec{Module: "enum"}
diff --git a/internal/printer/printer.go b/internal/printer/printer.go
@@ -130,6 +130,18 @@ func (w *writer) printNode(node *ast.Node, indent int32) {
 	}
 }
 
+func (w *writer) printAnnAssignWithMapped(aa *ast.AnnAssign, indent int32) {
+	if aa.Comment != "" {
+		w.print("# ")
+		w.print(aa.Comment)
+		w.print("\n")
+		w.printIndent(indent)
+	}
+	w.printName(aa.Target, indent)
+	w.print(": ")
+	w.printNode(aa.Annotation, indent)
+}
+
 func (w *writer) printAnnAssign(aa *ast.AnnAssign, indent int32) {
 	if aa.Comment != "" {
 		w.print("# ")
diff --git a/sqlc.yaml b/sqlc.yaml
@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://dist/sqlc-gen-python-orm.wasm
-    sha256: 2b6fe52d6ad5eb67cf783f5e57cd6cea3ce298115ca1ee7245215a41681dbdc5
+    sha256: 25d1a43b594104740b8c368956944bfe5482d6ada35a3b2faf7dfd968917c9c9
 sql:
 - schema: "_examples/schema.sql"
   queries: "_examples/query.sql"
