Fix #60

alembic_postgresql_enum/add_create_type_false.py

@@ -26,23 +26,34 @@ def __repr__(self):
         )
 
 
-def inject_repr_into_enums(column: Column):
-    """Swap postgresql.ENUM class to ReprWorkaround for the column type"""
-    if column.type.__class__ == sqlalchemy.Enum:
-        if not column.type.native_enum:
-            return
-        log.info("%r converted into postgresql.ENUM", column.type)
-        column.type = eval(repr(column.type).replace("Enum", "postgresql.ENUM"))
-    if isinstance(column.type, postgresql.ENUM):
-        if column.type.create_type:
-            log.info("create_type=False injected into %r", column.type.name)
-        replacement_enum_type = column.type
+log = logging.getLogger(f"alembic.{__name__}")
+
+
+def get_replacement_type(column_type):
+    replacement_enum_type = column_type
+
+    if replacement_enum_type.__class__ == sqlalchemy.Enum:
+        if not replacement_enum_type.native_enum:
+            return replacement_enum_type
+
+        log.info("%r converted into postgresql.ENUM", replacement_enum_type)
+        replacement_enum_type = eval(repr(replacement_enum_type).replace("Enum", "postgresql.ENUM"))
+
+    if isinstance(replacement_enum_type, postgresql.ENUM):
+        if replacement_enum_type.create_type:
+            log.info("create_type=False injected into %r", replacement_enum_type)
+
         replacement_enum_type.__class__ = ReprWorkaround
 
-        column.type = replacement_enum_type
+    return replacement_enum_type
 
 
-log = logging.getLogger(f"alembic.{__name__}")
+def inject_repr_into_enums(column: Column):
+    """Swap postgresql.ENUM class to ReprWorkaround for the column type"""
+    if isinstance(column.type, sqlalchemy.ARRAY):
+        column.type.item_type = get_replacement_type(column.type.item_type)
+    else:
+        column.type = get_replacement_type(column.type)
 
 
 def add_create_type_false(upgrade_ops: UpgradeOps):

tests/test_enum_creation/test_create_array.py

@@ -0,0 +1,58 @@
+from typing import TYPE_CHECKING, List
+
+import sqlalchemy
+
+from tests.base.render_and_run import compare_and_run
+from tests.schemas import (
+    get_car_schema_without_enum,
+    CAR_TABLE_NAME,
+    CAR_COLORS_COLUMN_NAME,
+    CAR_COLORS_ENUM_NAME,
+)
+
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+from sqlalchemy import Table, Column, Integer, MetaData
+
+
+def get_schema_with_enum_in_sqlalchemy_array_variants(variants: List[str]) -> MetaData:
+    schema = MetaData()
+
+    Table(
+        CAR_TABLE_NAME,
+        schema,
+        Column("id", Integer, primary_key=True),
+        Column(
+            CAR_COLORS_COLUMN_NAME,
+            sqlalchemy.ARRAY(sqlalchemy.Enum(*variants, name=CAR_COLORS_ENUM_NAME)),
+        ),
+    )
+
+    return schema
+
+
+def test_create_enum_on_create_table_with_array(connection: "Connection"):
+    """Check that library correctly creates enum before its use inside create_table. Enum is used in ARRAY"""
+    database_schema = get_car_schema_without_enum()
+    database_schema.create_all(connection)
+
+    new_enum_variants = ["black", "white", "red", "green", "blue", "other"]
+
+    target_schema = get_schema_with_enum_in_sqlalchemy_array_variants(new_enum_variants)
+
+    compare_and_run(
+        connection,
+        target_schema,
+        expected_upgrade=f"""
+        # ### commands auto generated by Alembic - please adjust! ###
+        sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind())
+        op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', sa.ARRAY(postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), nullable=True))
+        # ### end Alembic commands ###
+    """,
+        expected_downgrade=f"""
+        # ### commands auto generated by Alembic - please adjust! ###
+        op.drop_column('{CAR_TABLE_NAME}', '{CAR_COLORS_COLUMN_NAME}')
+        sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').drop(op.get_bind())
+        # ### end Alembic commands ###
+    """,
+    )

tests/test_enum_creation/test_create_table.py

@@ -65,7 +65,7 @@ def test_create_enum_on_create_table_with_array(connection: "Connection"):
         expected_upgrade=f"""
         # ### commands auto generated by Alembic - please adjust! ###
         sa.Enum({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind())
-        op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}')), nullable=True))
+        op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, new_enum_variants))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), nullable=True))
         # ### end Alembic commands ###
     """,
         expected_downgrade=f"""

tests/test_enum_creation/test_drop_column.py

@@ -79,7 +79,7 @@ def test_delete_enum_after_drop_column_with_array(connection: "Connection"):
         expected_downgrade=f"""
         # ### commands auto generated by Alembic - please adjust! ###
         sa.Enum({', '.join(map(repr, enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}').create(op.get_bind())
-        op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}')), autoincrement=False, nullable=True))
+        op.add_column('{CAR_TABLE_NAME}', sa.Column('{CAR_COLORS_COLUMN_NAME}', postgresql.ARRAY(postgresql.ENUM({', '.join(map(repr, enum_variants_to_delete))}, name='{CAR_COLORS_ENUM_NAME}', create_type=False)), autoincrement=False, nullable=True))
         # ### end Alembic commands ###
     """,
     )