From 961e19ad8775f766f8f83bef22d15e1717caeeb8 Mon Sep 17 00:00:00 2001 From: Dan Plischke Date: Tue, 19 Aug 2025 15:19:57 +0200 Subject: [PATCH 1/4] feat: handle SQLAlchemy types with missing python_type as typing.Any fix: SQLModel.metadata as metadataref --- CHANGES.rst | 4 ++++ src/sqlacodegen/generators.py | 8 +++++-- tests/test_generator_dataclass.py | 34 ++++++++++++++++++++++++++++- tests/test_generator_declarative.py | 33 +++++++++++++++++++++++++++- tests/test_generator_sqlmodel.py | 29 ++++++++++++++++++++++++ 5 files changed, 104 insertions(+), 4 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 96623d9a..8b7f19ec 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,6 +1,10 @@ Version history =============== +**UNRELEASED** +- Handle SQLAlchemy type with unimplemented python_type as typing.Any (PR by @danplischke) +- Fix SQLModel metadata reference (PR by @danplischke) + **3.1.0** - Type annotations for ARRAY column attributes now include the Python type of diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py index 7b4901a7..2eb3c6ea 100644 --- a/src/sqlacodegen/generators.py +++ b/src/sqlacodegen/generators.py @@ -1243,7 +1243,11 @@ def render_python_type(column_type: TypeEngine[Any]) -> str: if isinstance(column_type, DOMAIN): python_type = column_type.data_type.python_type else: - python_type = column_type.python_type + try: + python_type = column_type.python_type + except NotImplementedError: + self.add_literal_import("typing", "Any") + python_type = Any python_type_name = python_type.__name__ python_type_module = python_type.__module__ @@ -1435,7 +1439,7 @@ def generate_base(self) -> None: self.base = Base( literal_imports=[], declarations=[], - metadata_ref="", + metadata_ref="SQLModel.metadata", ) def collect_imports(self, models: Iterable[Model]) -> None: diff --git a/tests/test_generator_dataclass.py b/tests/test_generator_dataclass.py index 307f865c..afec050f 100644 --- a/tests/test_generator_dataclass.py +++ b/tests/test_generator_dataclass.py @@ -2,7 +2,7 @@ import pytest from _pytest.fixtures import FixtureRequest -from sqlalchemy.dialects.postgresql import UUID +from sqlalchemy.dialects.postgresql import TSVECTOR, UUID from sqlalchemy.engine import Engine from sqlalchemy.schema import Column, ForeignKeyConstraint, MetaData, Table from sqlalchemy.sql.expression import text @@ -267,3 +267,35 @@ class Simple(Base): id: Mapped[uuid.UUID] = mapped_column(UUID, primary_key=True) """, ) + + +def test_tsvector_missing_python_type(generator: CodeGenerator) -> None: + Table( + "simple_tsvector", + generator.metadata, + Column("id", UUID, primary_key=True), + Column("vector", TSVECTOR), + ) + + validate_code( + generator.generate(), + """\ + from typing import Any, Optional + import typing + import uuid + + from sqlalchemy import UUID + from sqlalchemy.dialects.postgresql import TSVECTOR + from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column + + class Base(MappedAsDataclass, DeclarativeBase): + pass + + + class SimpleTsvector(Base): + __tablename__ = 'simple_tsvector' + + id: Mapped[uuid.UUID] = mapped_column(UUID, primary_key=True) + vector: Mapped[Optional[typing.Any]] = mapped_column(TSVECTOR) + """, + ) diff --git a/tests/test_generator_declarative.py b/tests/test_generator_declarative.py index 931d5965..97b4a4eb 100644 --- a/tests/test_generator_declarative.py +++ b/tests/test_generator_declarative.py @@ -4,7 +4,7 @@ from _pytest.fixtures import FixtureRequest from sqlalchemy import BIGINT, PrimaryKeyConstraint from sqlalchemy.dialects import postgresql -from sqlalchemy.dialects.postgresql import JSON, JSONB +from sqlalchemy.dialects.postgresql import JSON, JSONB, TSVECTOR from sqlalchemy.engine import Engine from sqlalchemy.schema import ( CheckConstraint, @@ -1706,3 +1706,34 @@ class TestDomainJson(Base): foo: Mapped[Optional[dict]] = mapped_column(DOMAIN('domain_json', {domain_type.__name__}(astext_type=Text(length=128)), not_null=False)) """, ) + + +def test_tsvector_missing_python_type(generator: CodeGenerator) -> None: + Table( + "test_tsvector", + generator.metadata, + Column("id", BIGINT, primary_key=True), + Column("vector", TSVECTOR()), + ) + + validate_code( + generator.generate(), + """\ + from typing import Any, Optional + import typing + + from sqlalchemy import BigInteger + from sqlalchemy.dialects.postgresql import TSVECTOR + from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + + class Base(DeclarativeBase): + pass + + + class TestTsvector(Base): + __tablename__ = 'test_tsvector' + + id: Mapped[int] = mapped_column(BigInteger, primary_key=True) + vector: Mapped[Optional[typing.Any]] = mapped_column(TSVECTOR) + """, + ) diff --git a/tests/test_generator_sqlmodel.py b/tests/test_generator_sqlmodel.py index 32a736e2..4e88f73a 100644 --- a/tests/test_generator_sqlmodel.py +++ b/tests/test_generator_sqlmodel.py @@ -3,6 +3,7 @@ import pytest from _pytest.fixtures import FixtureRequest from sqlalchemy import Uuid +from sqlalchemy.dialects.postgresql import TSVECTOR from sqlalchemy.engine import Engine from sqlalchemy.schema import ( CheckConstraint, @@ -204,3 +205,31 @@ class SimpleUuid(SQLModel, table=True): id: uuid.UUID = Field(sa_column=Column('id', Uuid, primary_key=True)) """, ) + + +def test_tsvector_missing_python_type(generator: CodeGenerator) -> None: + Table( + "simple_tsvector", + generator.metadata, + Column("id", Uuid, primary_key=True), + Column("search", TSVECTOR), + ) + + validate_code( + generator.generate(), + """\ + from typing import Any, Optional + import typing + import uuid + + from sqlalchemy import Column, Uuid + from sqlalchemy.dialects.postgresql import TSVECTOR + from sqlmodel import Field, SQLModel + + class SimpleTsvector(SQLModel, table=True): + __tablename__ = 'simple_tsvector' + + id: uuid.UUID = Field(sa_column=Column('id', Uuid, primary_key=True)) + search: Optional[typing.Any] = Field(default=None, sa_column=Column('search', TSVECTOR)) + """, + ) From 74daca7492af2762cca414b90310af8d80d2dbef Mon Sep 17 00:00:00 2001 From: Dan Plischke Date: Tue, 19 Aug 2025 15:31:51 +0200 Subject: [PATCH 2/4] fix: python specialform type name in py3.9 --- src/sqlacodegen/generators.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py index 2eb3c6ea..5c2b2759 100644 --- a/src/sqlacodegen/generators.py +++ b/src/sqlacodegen/generators.py @@ -1249,7 +1249,11 @@ def render_python_type(column_type: TypeEngine[Any]) -> str: self.add_literal_import("typing", "Any") python_type = Any - python_type_name = python_type.__name__ + python_type_name = ( + python_type.__name__ + if hasattr(python_type, "__name__") + else python_type._name + ) python_type_module = python_type.__module__ if python_type_module == "builtins": return python_type_name From b2997f21de80cf82984be40d320309b224578199 Mon Sep 17 00:00:00 2001 From: Dan Plischke Date: Wed, 20 Aug 2025 09:56:15 +0200 Subject: [PATCH 3/4] integrate feedback --- CHANGES.rst | 1 + src/sqlacodegen/generators.py | 3 ++- tests/test_generator_sqlmodel.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/CHANGES.rst b/CHANGES.rst index 8b7f19ec..959634c0 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -2,6 +2,7 @@ Version history =============== **UNRELEASED** + - Handle SQLAlchemy type with unimplemented python_type as typing.Any (PR by @danplischke) - Fix SQLModel metadata reference (PR by @danplischke) diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py index 5c2b2759..628eb6c5 100644 --- a/src/sqlacodegen/generators.py +++ b/src/sqlacodegen/generators.py @@ -1246,7 +1246,7 @@ def render_python_type(column_type: TypeEngine[Any]) -> str: try: python_type = column_type.python_type except NotImplementedError: - self.add_literal_import("typing", "Any") + self.add_import(Any) python_type = Any python_type_name = ( @@ -1254,6 +1254,7 @@ def render_python_type(column_type: TypeEngine[Any]) -> str: if hasattr(python_type, "__name__") else python_type._name ) + python_type_module = python_type.__module__ if python_type_module == "builtins": return python_type_name diff --git a/tests/test_generator_sqlmodel.py b/tests/test_generator_sqlmodel.py index 4e88f73a..f5b72837 100644 --- a/tests/test_generator_sqlmodel.py +++ b/tests/test_generator_sqlmodel.py @@ -233,3 +233,33 @@ class SimpleTsvector(SQLModel, table=True): search: Optional[typing.Any] = Field(default=None, sa_column=Column('search', TSVECTOR)) """, ) + + +def test_metadata_ref(generator: CodeGenerator) -> None: + from sqlmodel import SQLModel + + Table( + "metadata_ref_test_table", + generator.metadata, + Column("id", INTEGER, primary_key=True), + ) + + code = generator.generate() + validate_code( + code, + """\ + from sqlalchemy import Column, Integer + from sqlmodel import Field, SQLModel + + class MetadataRefTestTable(SQLModel, table=True): + __tablename__ = 'metadata_ref_test_table' + + id: int = Field(sa_column=Column('id', Integer, primary_key=True)) + """, + ) + + SQLModel.metadata.clear() # clear the metadata to avoid with the tables defined in this test + exec(code, globals()) + + assert len(SQLModel.metadata.tables) == 1 + assert "metadata_ref_test_table" in SQLModel.metadata.tables From 6c79daa67212d70f63a2766cbd532735b42c7411 Mon Sep 17 00:00:00 2001 From: Dan Plischke Date: Wed, 20 Aug 2025 10:03:24 +0200 Subject: [PATCH 4/4] revert to using add_literal_import --- src/sqlacodegen/generators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py index 628eb6c5..db4fba0f 100644 --- a/src/sqlacodegen/generators.py +++ b/src/sqlacodegen/generators.py @@ -1246,7 +1246,7 @@ def render_python_type(column_type: TypeEngine[Any]) -> str: try: python_type = column_type.python_type except NotImplementedError: - self.add_import(Any) + self.add_literal_import("typing", "Any") python_type = Any python_type_name = (