Fixed same-name imports from wrong package (#411)

* Fixed same-name imports from wrong package

* Revert "Fixed same-name imports from wrong package"

This reverts commit b181bd7.

* Use render_python_column_type in SQLModelGenerator

* Use stdlib-list instead of sys.stdlib_module_names

* Update dependencies

* Use sys.stdlib_module_names for Python 3.10 or newer

* Use ClassVar instead of getter

* Use utility instead of ClassVar

* Fix indentation in pyproject.toml

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <[email protected]>

* PR Fixes

* Update CHANGES.rst

---------

Co-authored-by: Idan Sheinberg <[email protected]>
Co-authored-by: Alex Grönholm <[email protected]>

Author: 3 people

CHANGES.rst

@@ -12,6 +12,8 @@ Version history
 - Prevent double pluralization (PR by @dkratzert)
 - Fixes DOMAIN extending JSON/JSONB data types (PR by @sheinbergon)
 - Temporarily restrict SQLAlchemy version to 2.0.41 (PR by @sheinbergon)
+- Fixes ``add_import`` behavior when adding imports from sqlalchemy and overall better
+  alignment of import behavior(s) across generators
 
 **3.0.0**
 

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "SQLAlchemy >= 2.0.29,<2.0.42",
     "inflect >= 4.0.0",
     "importlib_metadata; python_version < '3.10'",
+    "stdlib-list; python_version < '3.10'"
 ]
 dynamic = ["version"]
 

src/sqlacodegen/generators.py

@@ -59,6 +59,7 @@
     get_common_fk_constraints,
     get_compiled_expression,
     get_constraint_sort_key,
+    get_stdlib_module_names,
     qualified_table_name,
     render_callable,
     uses_default_name,
@@ -119,9 +120,7 @@ def generate(self) -> str:
 @dataclass(eq=False)
 class TablesGenerator(CodeGenerator):
     valid_options: ClassVar[set[str]] = {"noindexes", "noconstraints", "nocomments"}
-    builtin_module_names: ClassVar[set[str]] = set(sys.builtin_module_names) | {
-        "dataclasses"
-    }
+    stdlib_module_names: ClassVar[set[str]] = get_stdlib_module_names()
 
     def __init__(
         self,
@@ -276,7 +275,7 @@ def add_import(self, obj: Any) -> None:
 
             if type_.__name__ in dialect_pkg.__all__:
                 pkgname = dialect_pkgname
-        elif type_.__name__ in dir(sqlalchemy):
+        elif type_ is getattr(sqlalchemy, type_.__name__, None):
             pkgname = "sqlalchemy"
         else:
             pkgname = type_.__module__
@@ -300,21 +299,26 @@ def group_imports(self) -> list[list[str]]:
         stdlib_imports: list[str] = []
         thirdparty_imports: list[str] = []
 
-        for package in sorted(self.imports):
-            imports = ", ".join(sorted(self.imports[package]))
+        def get_collection(package: str) -> list[str]:
             collection = thirdparty_imports
             if package == "__future__":
                 collection = future_imports
-            elif package in self.builtin_module_names:
+            elif package in self.stdlib_module_names:
                 collection = stdlib_imports
             elif package in sys.modules:
                 if "site-packages" not in (sys.modules[package].__file__ or ""):
                     collection = stdlib_imports
+            return collection
 
+        for package in sorted(self.imports):
+            imports = ", ".join(sorted(self.imports[package]))
+
+            collection = get_collection(package)
             collection.append(f"from {package} import {imports}")
 
         for module in sorted(self.module_imports):
-            thirdparty_imports.append(f"import {module}")
+            collection = get_collection(module)
+            collection.append(f"import {module}")
 
         return [
             group
@@ -1212,10 +1216,7 @@ def render_table_args(self, table: Table) -> str:
         else:
             return ""
 
-    def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
-        column = column_attr.column
-        rendered_column = self.render_column(column, column_attr.name != column.name)
-
+    def render_column_python_type(self, column: Column[Any]) -> str:
         def get_type_qualifiers() -> tuple[str, TypeEngine[Any], str]:
             column_type = column.type
             pre: list[str] = []
@@ -1254,7 +1255,14 @@ def render_python_type(column_type: TypeEngine[Any]) -> str:
 
         pre, col_type, post = get_type_qualifiers()
         column_python_type = f"{pre}{render_python_type(col_type)}{post}"
-        return f"{column_attr.name}: Mapped[{column_python_type}] = {rendered_column}"
+        return column_python_type
+
+    def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
+        column = column_attr.column
+        rendered_column = self.render_column(column, column_attr.name != column.name)
+        rendered_column_python_type = self.render_column_python_type(column)
+
+        return f"{column_attr.name}: Mapped[{rendered_column_python_type}] = {rendered_column}"
 
     def render_relationship(self, relationship: RelationshipAttribute) -> str:
         def render_column_attrs(column_attrs: list[ColumnAttribute]) -> str:
@@ -1444,15 +1452,6 @@ def collect_imports_for_model(self, model: Model) -> None:
             if model.relationships:
                 self.add_literal_import("sqlmodel", "Relationship")
 
-    def collect_imports_for_column(self, column: Column[Any]) -> None:
-        super().collect_imports_for_column(column)
-        try:
-            python_type = column.type.python_type
-        except NotImplementedError:
-            self.add_literal_import("typing", "Any")
-        else:
-            self.add_import(python_type)
-
     def render_module_variables(self, models: list[Model]) -> str:
         declarations: list[str] = []
         if any(not isinstance(model, ModelClass) for model in models):
@@ -1485,25 +1484,17 @@ def render_class_variables(self, model: ModelClass) -> str:
 
     def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
         column = column_attr.column
-        try:
-            python_type = column.type.python_type
-        except NotImplementedError:
-            python_type_name = "Any"
-        else:
-            python_type_name = python_type.__name__
+        rendered_column = self.render_column(column, True)
+        rendered_column_python_type = self.render_column_python_type(column)
 
         kwargs: dict[str, Any] = {}
-        if (
-            column.autoincrement and column.name in column.table.primary_key
-        ) or column.nullable:
-            self.add_literal_import("typing", "Optional")
+        if column.nullable:
             kwargs["default"] = None
-            python_type_name = f"Optional[{python_type_name}]"
-
-        rendered_column = self.render_column(column, True)
         kwargs["sa_column"] = f"{rendered_column}"
+
         rendered_field = render_callable("Field", kwargs=kwargs)
-        return f"{column_attr.name}: {python_type_name} = {rendered_field}"
+
+        return f"{column_attr.name}: {rendered_column_python_type} = {rendered_field}"
 
     def render_relationship(self, relationship: RelationshipAttribute) -> str:
         rendered = super().render_relationship(relationship).partition(" = ")[2]

src/sqlacodegen/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+import sys
 from collections.abc import Mapping
 from typing import Any, Literal, cast
 
@@ -206,3 +207,13 @@ def decode_postgresql_sequence(clause: TextClause) -> tuple[str | None, str | No
                 schema, sequence = sequence, ""
 
     return schema, sequence
+
+
+def get_stdlib_module_names() -> set[str]:
+    major, minor = sys.version_info.major, sys.version_info.minor
+    if (major, minor) > (3, 9):
+        return set(sys.builtin_module_names) | set(sys.stdlib_module_names)
+    else:
+        from stdlib_list import stdlib_list
+
+        return set(sys.builtin_module_names) | set(stdlib_list(f"{major}.{minor}"))

tests/test_cli.py

@@ -137,14 +137,11 @@ def test_cli_sqlmodels(db_path: Path, tmp_path: Path) -> None:
     assert (
         output_path.read_text()
         == """\
-from typing import Optional
-
 from sqlalchemy import Column, Integer, Text
 from sqlmodel import Field, SQLModel
 
 class Foo(SQLModel, table=True):
-    id: Optional[int] = Field(default=None, sa_column=Column('id', Integer, \
-primary_key=True))
+    id: int = Field(sa_column=Column('id', Integer, primary_key=True))
     name: str = Field(sa_column=Column('name', Text))
 """
     )

tests/test_generator_dataclass.py

@@ -251,10 +251,11 @@ def test_uuid_type_annotation(generator: CodeGenerator) -> None:
     validate_code(
         generator.generate(),
         """\
+        import uuid
+
         from sqlalchemy import UUID
         from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, \
 mapped_column
-        import uuid
 
         class Base(MappedAsDataclass, DeclarativeBase):
             pass

tests/test_generator_sqlmodel.py

@@ -2,6 +2,7 @@
 
 import pytest
 from _pytest.fixtures import FixtureRequest
+from sqlalchemy import Uuid
 from sqlalchemy.engine import Engine
 from sqlalchemy.schema import (
     CheckConstraint,
@@ -56,8 +57,7 @@ class Item(SQLModel, table=True):
                     Index('idx_text_number', 'text', 'number')
                 )
 
-                id: Optional[int] = Field(default=None, sa_column=Column(\
-'id', Integer, primary_key=True))
+                id: int = Field(sa_column=Column('id', Integer, primary_key=True))
                 number: Optional[int] = Field(default=None, sa_column=Column(\
 'number', Integer))
                 text: Optional[str] = Field(default=None, sa_column=Column(\
@@ -91,8 +91,7 @@ class SimpleConstraints(SQLModel, table=True):
                     UniqueConstraint('id', 'number')
                 )
 
-                id: Optional[int] = Field(default=None, sa_column=Column(\
-'id', Integer, primary_key=True))
+                id: int = Field(sa_column=Column('id', Integer, primary_key=True))
                 number: Optional[int] = Field(default=None, sa_column=Column(\
 'number', Integer))
         """,
@@ -124,8 +123,7 @@ def test_onetomany(generator: CodeGenerator) -> None:
             class SimpleContainers(SQLModel, table=True):
                 __tablename__ = 'simple_containers'
 
-                id: Optional[int] = Field(default=None, sa_column=Column(\
-'id', Integer, primary_key=True))
+                id: int = Field(sa_column=Column('id', Integer, primary_key=True))
 
                 simple_goods: list['SimpleGoods'] = Relationship(\
 back_populates='container')
@@ -134,8 +132,7 @@ class SimpleContainers(SQLModel, table=True):
             class SimpleGoods(SQLModel, table=True):
                 __tablename__ = 'simple_goods'
 
-                id: Optional[int] = Field(default=None, sa_column=Column(\
-'id', Integer, primary_key=True))
+                id: int = Field(sa_column=Column('id', Integer, primary_key=True))
                 container_id: Optional[int] = Field(default=None, sa_column=Column(\
 'container_id', ForeignKey('simple_containers.id')))
 
@@ -167,8 +164,7 @@ def test_onetoone(generator: CodeGenerator) -> None:
             class OtherItems(SQLModel, table=True):
                 __tablename__ = 'other_items'
 
-                id: Optional[int] = Field(default=None, sa_column=Column(\
-'id', Integer, primary_key=True))
+                id: int = Field(sa_column=Column('id', Integer, primary_key=True))
 
                 simple_onetoone: Optional['SimpleOnetoone'] = Relationship(\
 sa_relationship_kwargs={'uselist': False}, back_populates='other_item')
@@ -177,12 +173,34 @@ class OtherItems(SQLModel, table=True):
             class SimpleOnetoone(SQLModel, table=True):
                 __tablename__ = 'simple_onetoone'
 
-                id: Optional[int] = Field(default=None, sa_column=Column(\
-'id', Integer, primary_key=True))
+                id: int = Field(sa_column=Column('id', Integer, primary_key=True))
                 other_item_id: Optional[int] = Field(default=None, sa_column=Column(\
 'other_item_id', ForeignKey('other_items.id'), unique=True))
 
                 other_item: Optional['OtherItems'] = Relationship(\
 back_populates='simple_onetoone')
             """,
     )
+
+
+def test_uuid(generator: CodeGenerator) -> None:
+    Table(
+        "simple_uuid",
+        generator.metadata,
+        Column("id", Uuid, primary_key=True),
+    )
+
+    validate_code(
+        generator.generate(),
+        """\
+            import uuid
+
+            from sqlalchemy import Column, Uuid
+            from sqlmodel import Field, SQLModel
+
+            class SimpleUuid(SQLModel, table=True):
+                __tablename__ = 'simple_uuid'
+
+                id: uuid.UUID = Field(sa_column=Column('id', Uuid, primary_key=True))
+        """,
+    )