feat: add emit_str_enum config option (#66)

* feat: add emit_str_enum config option

* docs: add docs for emit_str_enum

* tests(emit_str_enum): add end to end test for emit str enum

* chore(tests): update wasm sha

* docs: add examples of with w/o

* docs: update to use correct dataclass syntax

README.md

@@ -1,21 +1,78 @@
 ## Usage
 
 ```yaml
-version: '2'
+version: "2"
 plugins:
-- name: py
-  wasm:
-    url: https://downloads.sqlc.dev/plugin/sqlc-gen-python_1.2.0.wasm
-    sha256: a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e
+  - name: py
+    wasm:
+      url: https://downloads.sqlc.dev/plugin/sqlc-gen-python_1.2.0.wasm
+      sha256: a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e
 sql:
-- schema: "schema.sql"
-  queries: "query.sql"
-  engine: postgresql
-  codegen:
-  - out: src/authors
-    plugin: py
-    options:
-      package: authors
-      emit_sync_querier: true
-      emit_async_querier: true
+  - schema: "schema.sql"
+    queries: "query.sql"
+    engine: postgresql
+    codegen:
+      - out: src/authors
+        plugin: py
+        options:
+          package: authors
+          emit_sync_querier: true
+          emit_async_querier: true
+```
+
+### Emit Pydantic Models instead of `dataclasses`
+
+Option: `emit_pydantic_models`
+
+By default, `sqlc-gen-python` will emit `dataclasses` for the models. If you prefer to use [`pydantic`](https://docs.pydantic.dev/latest/) models, you can enable this option.
+
+with `emit_pydantic_models`
+
+```py
+from pydantic import BaseModel
+
+class Author(pydantic.BaseModel):
+    id: int
+    name: str
+```
+
+without `emit_pydantic_models`
+
+```py
+import dataclasses
+
+@dataclasses.dataclass()
+class Author:
+    id: int
+    name: str
+```
+
+### Use `enum.StrEnum` for Enums
+
+Option: `emit_str_enum`
+
+`enum.StrEnum` was introduce in Python 3.11.
+
+`enum.StrEnum` is a subclass of `str` that is also a subclass of `Enum`. This allows for the use of `Enum` values as strings, compared to strings, or compared to other `enum.StrEnum` types.
+
+This is convenient for type checking and validation, as well as for serialization and deserialization.
+
+By default, `sqlc-gen-python` will emit `(str, enum.Enum)` for the enum classes. If you prefer to use `enum.StrEnum`, you can enable this option.
+
+with `emit_str_enum`
+
+```py
+class Status(enum.StrEnum):
+    """Venues can be either open or closed"""
+    OPEN = "op!en"
+    CLOSED = "clo@sed"
+```
+
+without `emit_str_enum` (current behavior)
+
+```py
+class Status(str, enum.Enum):
+    """Venues can be either open or closed"""
+    OPEN = "op!en"
+    CLOSED = "clo@sed"
 ```

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/endtoend/testdata/emit_pydantic_models/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/emit_str_enum/db/models.py

@@ -0,0 +1,19 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.27.0
+import dataclasses
+import enum
+from typing import Optional
+
+
+class BookStatus(enum.StrEnum):
+    AVAILABLE = "available"
+    CHECKED_OUT = "checked_out"
+    OVERDUE = "overdue"
+
+
+@dataclasses.dataclass()
+class Book:
+    id: int
+    title: str
+    status: Optional[BookStatus]

internal/endtoend/testdata/emit_str_enum/db/query.py

@@ -0,0 +1,111 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.27.0
+# source: query.sql
+from typing import AsyncIterator, Iterator, Optional
+
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+
+from db import models
+
+
+CREATE_BOOK = """-- name: create_book \\:one
+INSERT INTO books (
+          title, status
+) VALUES (
+  :p1, :p2
+) RETURNING id, title, status
+"""
+
+
+DELETE_BOOK = """-- name: delete_book \\:exec
+DELETE FROM books
+WHERE id = :p1
+"""
+
+
+GET_BOOK = """-- name: get_book \\:one
+SELECT id, title, status FROM books
+WHERE id = :p1 LIMIT 1
+"""
+
+
+LIST_BOOKS = """-- name: list_books \\:many
+SELECT id, title, status FROM books
+ORDER BY title
+"""
+
+
+class Querier:
+    def __init__(self, conn: sqlalchemy.engine.Connection):
+        self._conn = conn
+
+    def create_book(self, *, title: str, status: Optional[models.BookStatus]) -> Optional[models.Book]:
+        row = self._conn.execute(sqlalchemy.text(CREATE_BOOK), {"p1": title, "p2": status}).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            title=row[1],
+            status=row[2],
+        )
+
+    def delete_book(self, *, id: int) -> None:
+        self._conn.execute(sqlalchemy.text(DELETE_BOOK), {"p1": id})
+
+    def get_book(self, *, id: int) -> Optional[models.Book]:
+        row = self._conn.execute(sqlalchemy.text(GET_BOOK), {"p1": id}).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            title=row[1],
+            status=row[2],
+        )
+
+    def list_books(self) -> Iterator[models.Book]:
+        result = self._conn.execute(sqlalchemy.text(LIST_BOOKS))
+        for row in result:
+            yield models.Book(
+                id=row[0],
+                title=row[1],
+                status=row[2],
+            )
+
+
+class AsyncQuerier:
+    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
+        self._conn = conn
+
+    async def create_book(self, *, title: str, status: Optional[models.BookStatus]) -> Optional[models.Book]:
+        row = (await self._conn.execute(sqlalchemy.text(CREATE_BOOK), {"p1": title, "p2": status})).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            title=row[1],
+            status=row[2],
+        )
+
+    async def delete_book(self, *, id: int) -> None:
+        await self._conn.execute(sqlalchemy.text(DELETE_BOOK), {"p1": id})
+
+    async def get_book(self, *, id: int) -> Optional[models.Book]:
+        row = (await self._conn.execute(sqlalchemy.text(GET_BOOK), {"p1": id})).first()
+        if row is None:
+            return None
+        return models.Book(
+            id=row[0],
+            title=row[1],
+            status=row[2],
+        )
+
+    async def list_books(self) -> AsyncIterator[models.Book]:
+        result = await self._conn.stream(sqlalchemy.text(LIST_BOOKS))
+        async for row in result:
+            yield models.Book(
+                id=row[0],
+                title=row[1],
+                status=row[2],
+            )

internal/endtoend/testdata/emit_str_enum/query.sql

@@ -0,0 +1,18 @@
+-- name: GetBook :one
+SELECT * FROM books
+WHERE id = $1 LIMIT 1;
+
+-- name: ListBooks :many
+SELECT * FROM books
+ORDER BY title;
+
+-- name: CreateBook :one
+INSERT INTO books (
+          title, status
+) VALUES (
+  $1, $2
+) RETURNING *;
+
+-- name: DeleteBook :exec
+DELETE FROM books
+WHERE id = $1;

internal/endtoend/testdata/emit_str_enum/schema.sql

@@ -0,0 +1,8 @@
+CREATE TYPE book_status AS ENUM ('available', 'checked_out', 'overdue');
+
+
+CREATE TABLE books (
+          id     BIGSERIAL PRIMARY KEY,
+          title  text      NOT NULL,
+          status book_status DEFAULT 'available'
+);

internal/endtoend/testdata/emit_str_enum/sqlc.yaml

@@ -0,0 +1,19 @@
+version: "2"
+plugins:
+  - name: py
+    wasm:
+      url: file://../../../../bin/sqlc-gen-python.wasm
+      sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
+sql:
+  - schema: schema.sql
+    queries: query.sql
+    engine: postgresql
+    codegen:
+      - plugin: py
+        out: db
+        options:
+          package: db
+          emit_sync_querier: true
+          emit_async_querier: true
+          emit_str_enum: true
+

internal/endtoend/testdata/exec_result/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

internal/endtoend/testdata/exec_rows/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
 - name: py
   wasm:
     url: file://../../../../bin/sqlc-gen-python.wasm
-    sha256: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

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: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

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: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

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: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

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: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

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: "a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e"
+    sha256: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca"
 sql:
 - schema: schema.sql
   queries: query.sql

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{