Fixed AssertionError with uppercased type decorators (#375)

Fixes #315.

Author: agronholm

CHANGES.rst

@@ -9,6 +9,8 @@ Version history
 - Fixed two rendering issues in ``ENUM`` columns when a non-default schema is used: an
   unwarranted positional argument and missing the ``schema`` argument
 - Fixed ``AttributeError`` when metadata contains user defined column types
+- Fixed ``AssertionError`` when metadata contains a column type that is a type decorator
+  with an all-uppercase name
 
 **3.0.0rc5**
 

src/sqlacodegen/generators.py

@@ -36,6 +36,7 @@
     String,
     Table,
     Text,
+    TypeDecorator,
     UniqueConstraint,
 )
 from sqlalchemy.dialects.postgresql import JSONB
@@ -684,7 +685,7 @@ def get_adapted_type(self, coltype: Any) -> Any:
                 supercls, "__visit_name__"
             ):
                 # Don't try to adapt UserDefinedType as it's not a proper column type
-                if supercls is UserDefinedType:
+                if supercls is UserDefinedType or issubclass(supercls, TypeDecorator):
                     return coltype
 
                 # Hack to fix adaptation of the Enum class which is broken since

tests/test_generator_tables.py

@@ -4,6 +4,7 @@
 
 import pytest
 from _pytest.fixtures import FixtureRequest
+from sqlalchemy import TypeDecorator
 from sqlalchemy.dialects import mysql, postgresql
 from sqlalchemy.engine import Engine
 from sqlalchemy.schema import (
@@ -18,14 +19,19 @@
     UniqueConstraint,
 )
 from sqlalchemy.sql.expression import text
-from sqlalchemy.sql.sqltypes import NullType
+from sqlalchemy.sql.sqltypes import DateTime, NullType
 from sqlalchemy.types import INTEGER, NUMERIC, SMALLINT, VARCHAR, Text
 
 from sqlacodegen.generators import CodeGenerator, TablesGenerator
 
 from .conftest import validate_code
 
 
+# This needs to be uppercased to trigger #315
+class TIMESTAMP_DECORATOR(TypeDecorator[DateTime]):
+    impl = DateTime
+
+
 @pytest.fixture
 def generator(
     request: FixtureRequest, metadata: MetaData, engine: Engine
@@ -45,12 +51,15 @@ def test_fancy_coltypes(generator: CodeGenerator) -> None:
         Column("bool", postgresql.BOOLEAN),
         Column("vector", VECTOR(3)),
         Column("number", NUMERIC(10, asdecimal=False)),
+        Column("timestamp", TIMESTAMP_DECORATOR()),
         schema="someschema",
     )
 
     validate_code(
         generator.generate(),
         """\
+        from tests.test_generator_tables import TIMESTAMP_DECORATOR
+
         from pgvector.sqlalchemy.vector import VECTOR
         from sqlalchemy import Boolean, Column, Enum, MetaData, Numeric, Table
 
@@ -63,6 +72,7 @@ def test_fancy_coltypes(generator: CodeGenerator) -> None:
             Column('bool', Boolean),
             Column('vector', VECTOR(3)),
             Column('number', Numeric(10, asdecimal=False)),
+            Column('timestamp', TIMESTAMP_DECORATOR),
             schema='someschema'
         )
         """,