Fixed same-name imports from wrong package

Author: EthanKim8683

CHANGES.rst

@@ -11,6 +11,7 @@ Version history
   ``sqlalchemy-citext``, resolving ``PackageNotFoundError`` (PR by @oaimtiaz)
 - Prevent double pluralization (PR by @dkratzert)
 - Fixes DOMAIN extending JSON/JSONB data types (PR by @sheinbergon)
+- Fixes imports from wrong package if name shared
 
 **3.0.0**
 

src/sqlacodegen/generators.py

@@ -276,7 +276,9 @@ 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_.__name__ in dir(sqlalchemy) and type_ == getattr(
+            sqlalchemy, type_.__name__
+        ):
             pkgname = "sqlalchemy"
         else:
             pkgname = type_.__module__

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,
@@ -186,3 +187,27 @@ class SimpleOnetoone(SQLModel, table=True):
 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(),
+        """\
+            from typing import Optional
+            from uuid import UUID
+
+            from sqlalchemy import Column, Uuid
+            from sqlmodel import Field, SQLModel
+
+            class SimpleUuid(SQLModel, table=True):
+                __tablename__ = 'simple_uuid'
+
+                id: Optional[UUID] = Field(default=None, sa_column=Column('id', Uuid, primary_key=True))
+        """,
+    )