diff --git a/.gitignore b/.gitignore index 60b4f3d..c0a98a5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ -bin +bin/* +!bin/sqlc-gen-python.wasm # Devenv .envrc diff --git a/bin/sqlc-gen-python.wasm b/bin/sqlc-gen-python.wasm new file mode 100755 index 0000000..4a54a38 Binary files /dev/null and b/bin/sqlc-gen-python.wasm differ diff --git a/internal/config.go b/internal/config.go index 1a8a565..1746bdd 100644 --- a/internal/config.go +++ b/internal/config.go @@ -10,4 +10,11 @@ type Config struct { EmitStrEnum bool `json:"emit_str_enum"` QueryParameterLimit *int32 `json:"query_parameter_limit"` InflectionExcludeTableNames []string `json:"inflection_exclude_table_names"` + EmitModels *bool `json:"emit_models"` + ModelsPackage string `json:"models_package"` +} + +// ShouldEmitModels returns true unless emit_models was explicitly set to false. +func (c Config) ShouldEmitModels() bool { + return c.EmitModels == nil || *c.EmitModels } diff --git a/internal/endtoend/testdata/emit_models_package/db/query.py b/internal/endtoend/testdata/emit_models_package/db/query.py new file mode 100644 index 0000000..5992abb --- /dev/null +++ b/internal/endtoend/testdata/emit_models_package/db/query.py @@ -0,0 +1,73 @@ +# Code generated by sqlc. DO NOT EDIT. +# versions: +# sqlc v1.30.0 +# source: query.sql +from typing import Iterator, Optional + +import sqlalchemy + +from src.db.generated 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], + ) diff --git a/internal/endtoend/testdata/emit_models_package/query.sql b/internal/endtoend/testdata/emit_models_package/query.sql new file mode 100644 index 0000000..61142f7 --- /dev/null +++ b/internal/endtoend/testdata/emit_models_package/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; diff --git a/internal/endtoend/testdata/emit_models_package/schema.sql b/internal/endtoend/testdata/emit_models_package/schema.sql new file mode 100644 index 0000000..6842b27 --- /dev/null +++ b/internal/endtoend/testdata/emit_models_package/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' +); diff --git a/internal/endtoend/testdata/emit_models_package/sqlc.yaml b/internal/endtoend/testdata/emit_models_package/sqlc.yaml new file mode 100644 index 0000000..f553562 --- /dev/null +++ b/internal/endtoend/testdata/emit_models_package/sqlc.yaml @@ -0,0 +1,18 @@ +version: "2" +plugins: + - name: py + wasm: + url: file://../../../../bin/sqlc-gen-python.wasm + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" +sql: + - schema: schema.sql + queries: query.sql + engine: postgresql + codegen: + - plugin: py + out: db + options: + package: db + emit_sync_querier: true + emit_models: false + models_package: src.db.generated.models diff --git a/internal/endtoend/testdata/emit_pydantic_models/db/models.py b/internal/endtoend/testdata/emit_pydantic_models/db/models.py index 7676e5c..2300f77 100644 --- a/internal/endtoend/testdata/emit_pydantic_models/db/models.py +++ b/internal/endtoend/testdata/emit_pydantic_models/db/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import pydantic from typing import Optional diff --git a/internal/endtoend/testdata/emit_pydantic_models/db/query.py b/internal/endtoend/testdata/emit_pydantic_models/db/query.py index 6f5b76f..9835cb0 100644 --- a/internal/endtoend/testdata/emit_pydantic_models/db/query.py +++ b/internal/endtoend/testdata/emit_pydantic_models/db/query.py @@ -1,8 +1,8 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql -from typing import AsyncIterator, Iterator, Optional +from typing import Iterator, List, Optional import sqlalchemy import sqlalchemy.ext.asyncio @@ -102,11 +102,13 @@ async def get_author(self, *, id: int) -> Optional[models.Author]: bio=row[2], ) - async def list_authors(self) -> AsyncIterator[models.Author]: + async def list_authors(self) -> List[models.Author]: result = await self._conn.stream(sqlalchemy.text(LIST_AUTHORS)) + items = list() async for row in result: - yield models.Author( + items.append(models.Author( id=row[0], name=row[1], bio=row[2], - ) + )) + return items diff --git a/internal/endtoend/testdata/emit_pydantic_models/sqlc.yaml b/internal/endtoend/testdata/emit_pydantic_models/sqlc.yaml index beae200..5986929 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/emit_str_enum/db/models.py b/internal/endtoend/testdata/emit_str_enum/db/models.py index 5fdf754..5fd5508 100644 --- a/internal/endtoend/testdata/emit_str_enum/db/models.py +++ b/internal/endtoend/testdata/emit_str_enum/db/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses import enum from typing import Optional diff --git a/internal/endtoend/testdata/emit_str_enum/db/query.py b/internal/endtoend/testdata/emit_str_enum/db/query.py index 8082889..7c7e02b 100644 --- a/internal/endtoend/testdata/emit_str_enum/db/query.py +++ b/internal/endtoend/testdata/emit_str_enum/db/query.py @@ -1,8 +1,8 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql -from typing import AsyncIterator, Iterator, Optional +from typing import Iterator, List, Optional import sqlalchemy import sqlalchemy.ext.asyncio @@ -101,11 +101,13 @@ async def get_book(self, *, id: int) -> Optional[models.Book]: status=row[2], ) - async def list_books(self) -> AsyncIterator[models.Book]: + async def list_books(self) -> List[models.Book]: result = await self._conn.stream(sqlalchemy.text(LIST_BOOKS)) + items = list() async for row in result: - yield models.Book( + items.append(models.Book( id=row[0], title=row[1], status=row[2], - ) + )) + return items diff --git a/internal/endtoend/testdata/emit_str_enum/sqlc.yaml b/internal/endtoend/testdata/emit_str_enum/sqlc.yaml index 04e3feb..7597e07 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/exec_result/python/models.py b/internal/endtoend/testdata/exec_result/python/models.py index 034fb2d..ced3715 100644 --- a/internal/endtoend/testdata/exec_result/python/models.py +++ b/internal/endtoend/testdata/exec_result/python/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses diff --git a/internal/endtoend/testdata/exec_result/python/query.py b/internal/endtoend/testdata/exec_result/python/query.py index b68ce39..c063868 100644 --- a/internal/endtoend/testdata/exec_result/python/query.py +++ b/internal/endtoend/testdata/exec_result/python/query.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql import sqlalchemy import sqlalchemy.ext.asyncio diff --git a/internal/endtoend/testdata/exec_result/sqlc.yaml b/internal/endtoend/testdata/exec_result/sqlc.yaml index ddffc83..18e7d3f 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/exec_rows/python/models.py b/internal/endtoend/testdata/exec_rows/python/models.py index 034fb2d..ced3715 100644 --- a/internal/endtoend/testdata/exec_rows/python/models.py +++ b/internal/endtoend/testdata/exec_rows/python/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses diff --git a/internal/endtoend/testdata/exec_rows/python/query.py b/internal/endtoend/testdata/exec_rows/python/query.py index 7a9b2a6..c5a936d 100644 --- a/internal/endtoend/testdata/exec_rows/python/query.py +++ b/internal/endtoend/testdata/exec_rows/python/query.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql import sqlalchemy import sqlalchemy.ext.asyncio diff --git a/internal/endtoend/testdata/exec_rows/sqlc.yaml b/internal/endtoend/testdata/exec_rows/sqlc.yaml index ddffc83..18e7d3f 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/inflection_exclude_table_names/python/models.py b/internal/endtoend/testdata/inflection_exclude_table_names/python/models.py index 8ba8803..0614ac0 100644 --- a/internal/endtoend/testdata/inflection_exclude_table_names/python/models.py +++ b/internal/endtoend/testdata/inflection_exclude_table_names/python/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses diff --git a/internal/endtoend/testdata/inflection_exclude_table_names/python/query.py b/internal/endtoend/testdata/inflection_exclude_table_names/python/query.py index 1e1e161..8b9eb26 100644 --- a/internal/endtoend/testdata/inflection_exclude_table_names/python/query.py +++ b/internal/endtoend/testdata/inflection_exclude_table_names/python/query.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql from typing import Optional diff --git a/internal/endtoend/testdata/inflection_exclude_table_names/sqlc.yaml b/internal/endtoend/testdata/inflection_exclude_table_names/sqlc.yaml index efbb150..00ba68f 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/query_parameter_limit_two/python/models.py b/internal/endtoend/testdata/query_parameter_limit_two/python/models.py index 059675d..2ddf019 100644 --- a/internal/endtoend/testdata/query_parameter_limit_two/python/models.py +++ b/internal/endtoend/testdata/query_parameter_limit_two/python/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses diff --git a/internal/endtoend/testdata/query_parameter_limit_two/python/query.py b/internal/endtoend/testdata/query_parameter_limit_two/python/query.py index e8b723e..5a97c59 100644 --- a/internal/endtoend/testdata/query_parameter_limit_two/python/query.py +++ b/internal/endtoend/testdata/query_parameter_limit_two/python/query.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql import sqlalchemy import sqlalchemy.ext.asyncio diff --git a/internal/endtoend/testdata/query_parameter_limit_two/sqlc.yaml b/internal/endtoend/testdata/query_parameter_limit_two/sqlc.yaml index 336bca7..d42ad2b 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/query_parameter_limit_undefined/python/models.py b/internal/endtoend/testdata/query_parameter_limit_undefined/python/models.py index 30e80db..77bdfe5 100644 --- a/internal/endtoend/testdata/query_parameter_limit_undefined/python/models.py +++ b/internal/endtoend/testdata/query_parameter_limit_undefined/python/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses diff --git a/internal/endtoend/testdata/query_parameter_limit_undefined/python/query.py b/internal/endtoend/testdata/query_parameter_limit_undefined/python/query.py index 5a1fbbc..6380dce 100644 --- a/internal/endtoend/testdata/query_parameter_limit_undefined/python/query.py +++ b/internal/endtoend/testdata/query_parameter_limit_undefined/python/query.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql import sqlalchemy import sqlalchemy.ext.asyncio diff --git a/internal/endtoend/testdata/query_parameter_limit_undefined/sqlc.yaml b/internal/endtoend/testdata/query_parameter_limit_undefined/sqlc.yaml index c20cd57..5506780 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/endtoend/testdata/query_parameter_limit_zero/python/models.py b/internal/endtoend/testdata/query_parameter_limit_zero/python/models.py index 059675d..2ddf019 100644 --- a/internal/endtoend/testdata/query_parameter_limit_zero/python/models.py +++ b/internal/endtoend/testdata/query_parameter_limit_zero/python/models.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 import dataclasses diff --git a/internal/endtoend/testdata/query_parameter_limit_zero/python/query.py b/internal/endtoend/testdata/query_parameter_limit_zero/python/query.py index 47bd6a9..5edcd9c 100644 --- a/internal/endtoend/testdata/query_parameter_limit_zero/python/query.py +++ b/internal/endtoend/testdata/query_parameter_limit_zero/python/query.py @@ -1,6 +1,6 @@ # Code generated by sqlc. DO NOT EDIT. # versions: -# sqlc v1.28.0 +# sqlc v1.30.0 # source: query.sql import dataclasses diff --git a/internal/endtoend/testdata/query_parameter_limit_zero/sqlc.yaml b/internal/endtoend/testdata/query_parameter_limit_zero/sqlc.yaml index 6e2cdeb..c5be88b 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" 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 index c432e4f..44d94de 100644 --- 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: "d6846ffad948181e611e883cedd2d2be66e091edc1273a0abc6c9da18399e0ca" + sha256: "37ef74de82c70bcb6a96ea8e8c950c05e3054c5bfe3ea8a06b874d32fad20e5e" sql: - schema: schema.sql queries: query.sql diff --git a/internal/gen.go b/internal/gen.go index 6e50fae..c4dff0e 100644 --- a/internal/gen.go +++ b/internal/gen.go @@ -819,6 +819,19 @@ func asyncQuerierClassDef() *pyast.ClassDef { } } +// modelsImportModule returns the module to use in "from import models". +// When ModelsPackage is set (e.g. "src.db.generated.models"), returns its parent +// package ("src.db.generated"). Otherwise falls back to the configured Package. +func modelsImportModule(conf Config) string { + if conf.ModelsPackage != "" { + if idx := strings.LastIndex(conf.ModelsPackage, "."); idx >= 0 { + return conf.ModelsPackage[:idx] + } + return conf.ModelsPackage + } + return conf.Package +} + func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node { mod := moduleNode(ctx.SqlcVersion, source) std, pkg := i.queryImportSpecs(source) @@ -830,7 +843,7 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node { { Node: &pyast.Node_ImportFrom{ ImportFrom: &pyast.ImportFrom{ - Module: ctx.C.Package, + Module: modelsImportModule(ctx.C), Names: []*pyast.Node{ poet.Alias("models"), }, @@ -1023,21 +1036,32 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node { stream := connMethodNode("stream", q.ConstantName, q.ArgDictNode()) f.Body = append(f.Body, assignNode("result", poet.Await(stream)), + assignNode("items", poet.Node( + &pyast.Call{ + Func: poet.Name("list"), + }, + )), poet.Node( &pyast.AsyncFor{ Target: poet.Name("row"), Iter: poet.Name("result"), Body: []*pyast.Node{ poet.Expr( - poet.Yield( - q.Ret.RowNode("row"), + poet.Node( + &pyast.Call{ + Func: poet.Attribute(poet.Name("items"), "append"), + Args: []*pyast.Node{ + q.Ret.RowNode("row"), + }, + }, ), ), }, }, ), + poet.Return(poet.Name("items")), ) - f.Returns = subscriptNode("AsyncIterator", q.Ret.Annotation()) + f.Returns = subscriptNode("List", q.Ret.Annotation()) case ":exec": f.Body = append(f.Body, poet.Await(exec)) f.Returns = poet.Constant(nil) @@ -1112,9 +1136,10 @@ func Generate(_ context.Context, req *plugin.GenerateRequest) (*plugin.GenerateR } output := map[string]string{} - result := pyprint.Print(buildModelsTree(&tctx, i), pyprint.Options{}) - tctx.SourceName = "models.py" - output["models.py"] = string(result.Python) + if conf.ShouldEmitModels() { + result := pyprint.Print(buildModelsTree(&tctx, i), pyprint.Options{}) + output["models.py"] = string(result.Python) + } files := map[string]struct{}{} for _, q := range queries { diff --git a/internal/imports.go b/internal/imports.go index b88c58c..c1fec90 100644 --- a/internal/imports.go +++ b/internal/imports.go @@ -158,7 +158,7 @@ func (i *importer) queryImportSpecs(fileName string) (map[string]importSpec, map std["typing.Iterator"] = importSpec{Module: "typing", Name: "Iterator"} } if i.C.EmitAsyncQuerier { - std["typing.AsyncIterator"] = importSpec{Module: "typing", Name: "AsyncIterator"} + std["typing.List"] = importSpec{Module: "typing", Name: "List"} } } queryValueModelImports(q.Ret) @@ -173,7 +173,8 @@ func (i *importer) queryImportSpecs(fileName string) (map[string]importSpec, map func (i *importer) queryImports(fileName string) []string { std, pkg := i.queryImportSpecs(fileName) - modelImportStr := fmt.Sprintf("from %s import models", i.C.Package) + importMod := modelsImportModule(i.C) + modelImportStr := fmt.Sprintf("from %s import models", importMod) importLines := []string{ buildImportBlock(std),