Merge pull request #90 from dswistowski/develop

feature: fixes #33 allow to ignore enums by name

README.md

@@ -52,6 +52,7 @@ alembic_postgresql_enum.set_configuration(
 * [Creation of new enum values](#creation-of-new-enum-values)
 * [Deletion of enums values](#deletion-of-enums-values)
 * [Renaming of enum values](#rename-enum-value)
+* [Omitting managing enums](#omitting-managing-enums)
 
 ### Creation of enum<a id="creation-of-enum"></a>
 
@@ -284,3 +285,21 @@ def downgrade():
 Do not forget to switch places old and new values for downgrade
 
 All defaults in postgres will be renamed automatically as well
+
+### Omitting managing enums<a id="omitting-managing-enums"></a>
+
+If configured `include_name` function returns `False` given enum will be not managed.
+```python
+import alembic_postgresql_enum
+
+def include_name(name: str) -> bool:
+    return name not in ['enum-to-ignore', 'some-internal-enum']
+
+alembic_postgresql_enum.set_configuration(
+    alembic_postgresql_enum.Config(
+        include_name=include_name,
+    )
+)
+```
+
+Feature is similar to [sqlalchemy feature for tables](https://alembic.sqlalchemy.org/en/latest/autogenerate.html#omitting-table-names-from-the-autogenerate-process)

alembic_postgresql_enum/compare_dispatch.py

@@ -15,7 +15,7 @@
     drop_unused_enums,
 )
 from alembic_postgresql_enum.get_enum_data import get_defined_enums, get_declared_enums
-
+from alembic_postgresql_enum.configuration import get_configuration
 
 log = logging.getLogger(f"alembic.{__name__}")
 
@@ -48,6 +48,8 @@ def compare_enums(
     add_create_type_false(upgrade_ops)
     add_postgres_using_to_text(upgrade_ops)
 
+    configuration = get_configuration()
+
     schema_names = list(schema_names)
 
     # Issue #40
@@ -61,9 +63,15 @@ def compare_enums(
         if schema is None:
             schema = default_schema
 
-        definitions = get_defined_enums(autogen_context.connection, schema)
+        definitions = get_defined_enums(autogen_context.connection, schema, configuration.include_name)
+
         declarations = get_declared_enums(
-            autogen_context.metadata, schema, default_schema, autogen_context.connection, upgrade_ops
+            autogen_context.metadata,
+            schema,
+            default_schema,
+            autogen_context.connection,
+            upgrade_ops,
+            configuration.include_name,
         )
 
         create_new_enums(definitions, declarations.enum_values, schema, upgrade_ops)

alembic_postgresql_enum/configuration.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass
+from typing import Callable
 
 
 @dataclass
 class Config:
     add_type_ignore: bool = False
+    include_name: Callable[[str], bool] = lambda _: True
 
 
 _config = Config()

alembic_postgresql_enum/get_enum_data/declared_enums.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Tuple, Any, Set, Union, List, TYPE_CHECKING, cast, Optional
+from typing import Tuple, Any, Set, Union, List, TYPE_CHECKING, cast, Optional, Callable
 
 import sqlalchemy
 from alembic.operations.ops import UpgradeOps
@@ -52,6 +52,7 @@ def get_declared_enums(
     default_schema: str,
     connection: "Connection",
     upgrade_ops: Optional[UpgradeOps] = None,
+    include_name: Callable[[str], bool] = lambda _: True,
 ) -> DeclaredEnumValues:
     """
     Return a dict mapping SQLAlchemy declared enumeration types to the set of their values
@@ -100,6 +101,9 @@ def get_declared_enums(
                 if not column_type_is_enum(column_type):
                     continue
 
+                if not include_name(column_type.name):
+                    continue
+
                 column_type_schema = column_type.schema or default_schema  # type: ignore[attr-defined]
                 if column_type_schema != schema:
                     continue

alembic_postgresql_enum/get_enum_data/defined_enums.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 from alembic_postgresql_enum.get_enum_data.types import EnumNamesToValues
 from alembic_postgresql_enum.sql_commands.enum_type import get_all_enums
@@ -16,17 +16,28 @@ def _remove_schema_prefix(enum_name: str, schema: str) -> str:
     return enum_name
 
 
-def get_defined_enums(connection: "Connection", schema: str) -> EnumNamesToValues:
+def get_defined_enums(
+    connection: "Connection", schema: str, include_name: Callable[[str], bool] = lambda _: True
+) -> EnumNamesToValues:
     """
     Return a dict mapping PostgreSQL defined enumeration types to the set of their
     defined values.
     :param conn:
         SQLAlchemy connection instance.
     :param str schema:
         Schema name (e.g. "public").
+    :param Callable[[str], bool] include_name:
+        Callable that returns True if the enum name should be included
     :returns DeclaredEnumValues:
         enum_definitions={
             "my_enum": tuple(["a", "b", "c"]),
         }
     """
-    return {_remove_schema_prefix(name, schema): tuple(values) for name, values in get_all_enums(connection, schema)}
+
+    return {
+        enum_name: tuple(values)
+        for enum_name, values in (
+            (_remove_schema_prefix(name, schema), values) for name, values in get_all_enums(connection, schema)
+        )
+        if include_name(enum_name)
+    }

tests/get_enum_data/test_get_declared_enums.py

@@ -1,3 +1,4 @@
+import random
 from typing import TYPE_CHECKING
 
 from alembic_postgresql_enum.get_enum_data import TableReference, get_declared_enums
@@ -40,3 +41,27 @@ def test_with_multiple_enums(connection: "Connection"):
     function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection)
 
     assert function_result == declared_enum_values
+
+
+def test_can_filter_by_name(connection: "Connection"):
+    declared_enum_values = get_declared_enum_values_with_orders_and_users()
+    declared_schema = get_schema_by_declared_enum_values(declared_enum_values)
+
+    all_declared_enums = declared_enum_values.enum_values
+    enum_to_exclude = random.choice(list(all_declared_enums))
+    enum_to_include = [e for e in all_declared_enums if e != enum_to_exclude]
+
+    function_result = get_declared_enums(
+        declared_schema,
+        DEFAULT_SCHEMA,
+        DEFAULT_SCHEMA,
+        connection,
+        None,
+        lambda enum_name: enum_name not in [enum_to_exclude],
+    )
+
+    assert enum_to_exclude not in function_result.enum_values
+    assert enum_to_exclude not in function_result.enum_table_references
+
+    assert set(enum_to_include) == set(function_result.enum_values)
+    assert set(enum_to_include) == set(function_result.enum_table_references)

tests/get_enum_data/test_get_defined_enums.py

@@ -1,3 +1,4 @@
+import random
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -31,3 +32,17 @@ def test_with_multiple_enums(connection: "Connection"):
     function_result = get_defined_enums(connection, DEFAULT_SCHEMA)
 
     assert function_result == declared_enum_values.enum_values
+
+
+def test_get_defined_enums_will_filter_by_name(connection: "Connection"):
+    declared_enum_values = get_declared_enum_values_with_orders_and_users()
+    defined_schema = get_schema_by_declared_enum_values(declared_enum_values)
+    all_declared_enums = declared_enum_values.enum_values
+    enum_to_exclude = random.choice(list(all_declared_enums))
+
+    defined_schema.create_all(connection)
+
+    function_result = get_defined_enums(
+        connection, DEFAULT_SCHEMA, lambda enum_name: enum_name not in [enum_to_exclude]
+    )
+    assert enum_to_exclude not in function_result