From 4fc55fb62906f49c3b365492d64051629db0d80e Mon Sep 17 00:00:00 2001 From: "artem.golovin" Date: Sat, 16 Nov 2024 22:52:31 +0400 Subject: [PATCH] Hide type: ignore behind the `add_type_ignore` feature flag --- README.md | 97 ++++++++++++++----- alembic_postgresql_enum/__init__.py | 1 + alembic_postgresql_enum/configuration.py | 18 ++++ .../operations/sync_enum_values.py | 12 ++- tests/base/run_migration_test_abc.py | 6 ++ tests/sync_enum_values/test_array_column.py | 74 +++++++------- tests/sync_enum_values/test_render.py | 81 ++++++++++------ .../test_run_array_new_column.py | 41 ++++---- 8 files changed, 220 insertions(+), 110 deletions(-) create mode 100644 alembic_postgresql_enum/configuration.py diff --git a/README.md b/README.md index 89cba9b..ea60f8f 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +from alembic_postgresql_enum.configuration import Config + # alembic-postgresql-enum [](https://pypi.org/project/alembic-postgresql-enum/) [](https://pypi.org/project/alembic-postgresql-enum/) @@ -30,6 +32,21 @@ import alembic_postgresql_enum To the top of your migrations/env.py file. +## Configuration + +You can configure this extension to disable parts of it, or to enable some feature flags + +To do so you need to call set_configuration function after the import: + +```python +import alembic_postgresql_enum + +alembic_postgresql_enum.set_configuration( + alembic_postgresql_enum.Config( + add_type_ignore=True, + ) +) +``` ## Features * [Creation of enums](#creation-of-enum) @@ -147,17 +164,25 @@ class MyEnum(enum.Enum): ```python def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'myenum', ['one', 'two', 'three', 'four'], - [('example_table', 'enum_field')], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'three', 'four'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[], + ) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'myenum', ['one', 'two', 'three'], - [('example_table', 'enum_field')], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'three'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[], + ) # ### end Alembic commands ### ``` @@ -175,17 +200,25 @@ class MyEnum(enum.Enum): ```python def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'myenum', ['one', 'two'], - [('example_table', 'enum_field')], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[], + ) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'myenum', ['one', 'two', 'three'], - [('example_table', 'enum_field')], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'three'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[], + ) # ### end Alembic commands ### ``` @@ -203,17 +236,25 @@ This code will generate this migration: ```python def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'myenum', ['one', 'two', 'three'], - [('example_table', 'enum_field')], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'three'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[], + ) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'myenum', ['one', 'two', 'tree'], - [('example_table', 'enum_field')], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'tree'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[], + ) # ### end Alembic commands ### ``` @@ -223,15 +264,23 @@ So adjust migration like that ```python def upgrade(): - op.sync_enum_values('public', 'myenum', ['one', 'two', 'three'], - [('example_table', 'enum_field')], - enum_values_to_rename=[('tree', 'three')]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'three'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[('tree', 'three')], + ) def downgrade(): - op.sync_enum_values('public', 'myenum', ['one', 'two', 'tree'], - [('example_table', 'enum_field')], - enum_values_to_rename=[('three', 'tree')]) + op.sync_enum_values( + enum_schema='public', + enum_name='myenum', + new_values=['one', 'two', 'tree'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')], + enum_values_to_rename=[('three', 'tree')], + ) ``` Do not forget to switch places old and new values for downgrade diff --git a/alembic_postgresql_enum/__init__.py b/alembic_postgresql_enum/__init__.py index 6ea9bbb..8b62712 100644 --- a/alembic_postgresql_enum/__init__.py +++ b/alembic_postgresql_enum/__init__.py @@ -1,2 +1,3 @@ from .compare_dispatch import compare_enums from .get_enum_data import ColumnType, TableReference +from .configuration import set_configuration, Config diff --git a/alembic_postgresql_enum/configuration.py b/alembic_postgresql_enum/configuration.py new file mode 100644 index 0000000..29eceed --- /dev/null +++ b/alembic_postgresql_enum/configuration.py @@ -0,0 +1,18 @@ +from dataclasses import dataclass + + +@dataclass +class Config: + add_type_ignore: bool = False + + +_config = Config() + + +def set_configuration(config: Config): + global _config + _config = config + + +def get_configuration() -> Config: + return _config diff --git a/alembic_postgresql_enum/operations/sync_enum_values.py b/alembic_postgresql_enum/operations/sync_enum_values.py index 2efcd6d..a642107 100644 --- a/alembic_postgresql_enum/operations/sync_enum_values.py +++ b/alembic_postgresql_enum/operations/sync_enum_values.py @@ -7,6 +7,7 @@ from alembic.autogenerate.api import AutogenContext from sqlalchemy.exc import DataError +from alembic_postgresql_enum.configuration import get_configuration from alembic_postgresql_enum.get_enum_data.types import Unspecified from alembic_postgresql_enum.sql_commands.column_default import ( get_column_default, @@ -203,12 +204,17 @@ def is_column_type_import_needed(self) -> bool: @alembic.autogenerate.render.renderers.dispatch_for(SyncEnumValuesOp) def render_sync_enum_value_op(autogen_context: AutogenContext, op: SyncEnumValuesOp): + config = get_configuration() if op.is_column_type_import_needed: autogen_context.imports.add("from alembic_postgresql_enum import ColumnType") autogen_context.imports.add("from alembic_postgresql_enum import TableReference") return ( - f"op.sync_enum_values({op.schema!r}, {op.name!r}, {op.new_values!r}, # type: ignore[attr-defined]\n" - f" {op.affected_columns!r},\n" - f" enum_values_to_rename=[])" + f"op.sync_enum_values({' # type: ignore[attr-defined]' if config.add_type_ignore else ''}\n" + f" enum_schema={op.schema!r},\n" + f" enum_name={op.name!r},\n" + f" new_values={op.new_values!r},\n" + f" affected_columns={op.affected_columns!r},\n" + f" enum_values_to_rename=[],\n" + f")" ) diff --git a/tests/base/run_migration_test_abc.py b/tests/base/run_migration_test_abc.py index eb6069a..64dfb76 100644 --- a/tests/base/run_migration_test_abc.py +++ b/tests/base/run_migration_test_abc.py @@ -1,6 +1,8 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING +import alembic_postgresql_enum +from alembic_postgresql_enum.configuration import Config, get_configuration from tests.base.render_and_run import compare_and_run if TYPE_CHECKING: @@ -14,6 +16,7 @@ class CompareAndRunTestCase(ABC): """ disable_running = False + config = Config() @abstractmethod def get_database_schema(self) -> MetaData: ... @@ -31,6 +34,8 @@ def get_expected_upgrade(self) -> str: ... def get_expected_downgrade(self) -> str: ... def test_run(self, connection: "Connection"): + old_config = get_configuration() + alembic_postgresql_enum.set_configuration(self.config) database_schema = self.get_database_schema() target_schema = self.get_target_schema() @@ -44,3 +49,4 @@ def test_run(self, connection: "Connection"): expected_downgrade=self.get_expected_downgrade(), disable_running=self.disable_running, ) + alembic_postgresql_enum.set_configuration(old_config) diff --git a/tests/sync_enum_values/test_array_column.py b/tests/sync_enum_values/test_array_column.py index 3800998..6bf7542 100644 --- a/tests/sync_enum_values/test_array_column.py +++ b/tests/sync_enum_values/test_array_column.py @@ -7,9 +7,11 @@ from alembic_postgresql_enum import ColumnType from alembic_postgresql_enum.get_enum_data import TableReference from alembic_postgresql_enum.operations import SyncEnumValuesOp +from tests.base.run_migration_test_abc import CompareAndRunTestCase if TYPE_CHECKING: from sqlalchemy import Connection +from sqlalchemy import MetaData from tests.schemas import ( get_schema_with_enum_in_array_variants, @@ -21,43 +23,47 @@ from tests.utils.migration_context import create_migration_context -def test_add_new_enum_value_render_with_array(connection: "Connection"): +class TestAddNewEnumValueRenderWithArray(CompareAndRunTestCase): """Check that enum variants are updated when new variant is added""" - old_enum_variants = ["black", "white", "red", "green", "blue", "other"] - - database_schema = get_schema_with_enum_in_array_variants(old_enum_variants) - database_schema.create_all(connection) - - new_enum_variants = old_enum_variants.copy() - new_enum_variants.append("violet") - - target_schema = get_schema_with_enum_in_array_variants(new_enum_variants) - context = create_migration_context(connection, target_schema) - - template_args = {} - autogenerate._render_migration_diffs(context, template_args) - - assert template_args["imports"] == ( - "from alembic_postgresql_enum import ColumnType" "\nfrom alembic_postgresql_enum import TableReference" - ) + old_enum_variants = ["black", "white", "red", "green", "blue", "other"] + new_enum_variants = old_enum_variants + ["violet"] + + def get_database_schema(self) -> MetaData: + schema = get_schema_with_enum_in_array_variants(self.old_enum_variants) + return schema + + def get_target_schema(self) -> MetaData: + schema = get_schema_with_enum_in_array_variants(self.new_enum_variants) + return schema + + def get_expected_upgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='{DEFAULT_SCHEMA}', + enum_name='{CAR_COLORS_ENUM_NAME}', + new_values=[{', '.join(map(repr, self.new_enum_variants))}], + affected_columns=[TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{CAR_TABLE_NAME}', column_name='{CAR_COLORS_COLUMN_NAME}', column_type=ColumnType.ARRAY)], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ + + def get_expected_downgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='{DEFAULT_SCHEMA}', + enum_name='{CAR_COLORS_ENUM_NAME}', + new_values=[{', '.join(map(repr, self.old_enum_variants))}], + affected_columns=[TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{CAR_TABLE_NAME}', column_name='{CAR_COLORS_COLUMN_NAME}', column_type=ColumnType.ARRAY)], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ - assert ( - template_args["upgrades"] - == f"""# ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('{DEFAULT_SCHEMA}', '{CAR_COLORS_ENUM_NAME}', [{', '.join(map(repr, new_enum_variants))}], # type: ignore[attr-defined] - [TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{CAR_TABLE_NAME}', column_name='{CAR_COLORS_COLUMN_NAME}', column_type=ColumnType.ARRAY)], - enum_values_to_rename=[]) - # ### end Alembic commands ###""" - ) - assert ( - template_args["downgrades"] - == f"""# ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('{DEFAULT_SCHEMA}', '{CAR_COLORS_ENUM_NAME}', [{', '.join(map(repr, old_enum_variants))}], # type: ignore[attr-defined] - [TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{CAR_TABLE_NAME}', column_name='{CAR_COLORS_COLUMN_NAME}', column_type=ColumnType.ARRAY)], - enum_values_to_rename=[]) - # ### end Alembic commands ###""" - ) + "from alembic_postgresql_enum import ColumnType" "\nfrom alembic_postgresql_enum import TableReference" def test_add_new_enum_value_diff_tuple_with_array(connection: "Connection"): diff --git a/tests/sync_enum_values/test_render.py b/tests/sync_enum_values/test_render.py index 3acf069..c85eb8c 100644 --- a/tests/sync_enum_values/test_render.py +++ b/tests/sync_enum_values/test_render.py @@ -7,6 +7,7 @@ from sqlalchemy.dialects import postgresql from sqlalchemy.orm import declarative_base +from alembic_postgresql_enum.configuration import Config from alembic_postgresql_enum.get_enum_data import TableReference from alembic_postgresql_enum.operations import SyncEnumValuesOp from tests.base.run_migration_test_abc import CompareAndRunTestCase @@ -25,39 +26,47 @@ from tests.utils.migration_context import create_migration_context -def test_add_new_enum_value_render(connection: "Connection"): +class TestAddNewEnumValueRender(CompareAndRunTestCase): """Check that enum variants are updated when new variant is added""" - old_enum_variants = ["active", "passive"] - database_schema = get_schema_with_enum_variants(old_enum_variants) - database_schema.create_all(connection) + config = Config(add_type_ignore=True) - new_enum_variants = old_enum_variants.copy() - new_enum_variants.append("banned") + old_enum_variants = ["active", "passive"] + new_enum_variants = old_enum_variants + ["banned"] - target_schema = get_schema_with_enum_variants(new_enum_variants) + def get_database_schema(self) -> MetaData: + schema = get_schema_with_enum_variants(self.old_enum_variants) + return schema - context = create_migration_context(connection, target_schema) + def get_target_schema(self) -> MetaData: + schema = get_schema_with_enum_variants(self.new_enum_variants) + return schema - template_args = {} - autogenerate._render_migration_diffs(context, template_args) + def get_expected_upgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( # type: ignore[attr-defined] + enum_schema='{DEFAULT_SCHEMA}', + enum_name='{USER_STATUS_ENUM_NAME}', + new_values=[{', '.join(map(repr, self.new_enum_variants))}], + affected_columns=[TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{USER_TABLE_NAME}', column_name='{USER_STATUS_COLUMN_NAME}')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ - assert ( - template_args["upgrades"] - == f"""# ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, new_enum_variants))}], # type: ignore[attr-defined] - [TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{USER_TABLE_NAME}', column_name='{USER_STATUS_COLUMN_NAME}')], - enum_values_to_rename=[]) - # ### end Alembic commands ###""" - ) - assert ( - template_args["downgrades"] - == f"""# ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, old_enum_variants))}], # type: ignore[attr-defined] - [TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{USER_TABLE_NAME}', column_name='{USER_STATUS_COLUMN_NAME}')], - enum_values_to_rename=[]) - # ### end Alembic commands ###""" - ) + def get_expected_downgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( # type: ignore[attr-defined] + enum_schema='{DEFAULT_SCHEMA}', + enum_name='{USER_STATUS_ENUM_NAME}', + new_values=[{', '.join(map(repr, self.old_enum_variants))}], + affected_columns=[TableReference(table_schema='{DEFAULT_SCHEMA}', table_name='{USER_TABLE_NAME}', column_name='{USER_STATUS_COLUMN_NAME}')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ def test_add_new_enum_value_diff_tuple(connection: "Connection"): @@ -207,17 +216,25 @@ class ExampleTable(Base): def get_expected_upgrade(self) -> str: return """ # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'my_enum', ['one', 'two', 'three', 'four'], # type: ignore[attr-defined] - [TableReference(table_schema='public', table_name='example_table', column_name='enum_field', existing_server_default="'one'::my_enum")], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='my_enum', + new_values=['one', 'two', 'three', 'four'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field', existing_server_default="'one'::my_enum")], + enum_values_to_rename=[], + ) # ### end Alembic commands ### """ def get_expected_downgrade(self) -> str: return """ # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'my_enum', ['one', 'two', 'three'], # type: ignore[attr-defined] - [TableReference(table_schema='public', table_name='example_table', column_name='enum_field', existing_server_default="'one'::my_enum")], - enum_values_to_rename=[]) + op.sync_enum_values( + enum_schema='public', + enum_name='my_enum', + new_values=['one', 'two', 'three'], + affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field', existing_server_default="'one'::my_enum")], + enum_values_to_rename=[], + ) # ### end Alembic commands ### """ diff --git a/tests/sync_enum_values/test_run_array_new_column.py b/tests/sync_enum_values/test_run_array_new_column.py index 00f2aa3..50e9b02 100644 --- a/tests/sync_enum_values/test_run_array_new_column.py +++ b/tests/sync_enum_values/test_run_array_new_column.py @@ -1,15 +1,12 @@ from enum import Enum -from typing import TYPE_CHECKING import sqlalchemy -from sqlalchemy import MetaData, Table, Column, insert +from sqlalchemy import MetaData, Table, Column from sqlalchemy.dialects import postgresql +from alembic_postgresql_enum.configuration import Config from tests.base.run_migration_test_abc import CompareAndRunTestCase -if TYPE_CHECKING: - from sqlalchemy import Connection - class OldEnum(Enum): A = "a" @@ -23,6 +20,8 @@ class NewEnum(Enum): class TestNewArrayColumnColumn(CompareAndRunTestCase): + config = Config(add_type_ignore=True) + def get_database_schema(self) -> MetaData: database_schema = MetaData() Table("a", database_schema) # , Column("value", postgresql.ARRAY(postgresql.ENUM(OldEnum))) @@ -61,20 +60,28 @@ def get_target_schema(self) -> MetaData: def get_expected_upgrade(self) -> str: return """ - # ### commands auto generated by Alembic - please adjust! ### - op.add_column('a', sa.Column('value', postgresql.ARRAY(postgresql.ENUM('A', 'B', 'C', name='my_enum', create_type=False)), server_default=sa.text("ARRAY['A', 'B']::my_enum[]"), nullable=True)) - op.sync_enum_values('public', 'my_enum', ['A', 'B', 'C'], # type: ignore[attr-defined] - [TableReference(table_schema='public', table_name='a', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A', 'B']::my_enum[]"), TableReference(table_schema='public', table_name='b', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A'::my_enum, 'B'::my_enum]")], - enum_values_to_rename=[]) - # ### end Alembic commands ### + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('a', sa.Column('value', postgresql.ARRAY(postgresql.ENUM('A', 'B', 'C', name='my_enum', create_type=False)), server_default=sa.text("ARRAY['A', 'B']::my_enum[]"), nullable=True)) + op.sync_enum_values( # type: ignore[attr-defined] + enum_schema='public', + enum_name='my_enum', + new_values=['A', 'B', 'C'], + affected_columns=[TableReference(table_schema='public', table_name='a', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A', 'B']::my_enum[]"), TableReference(table_schema='public', table_name='b', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A'::my_enum, 'B'::my_enum]")], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### """ def get_expected_downgrade(self) -> str: return """ - # ### commands auto generated by Alembic - please adjust! ### - op.sync_enum_values('public', 'my_enum', ['A', 'B'], # type: ignore[attr-defined] - [TableReference(table_schema='public', table_name='a', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A', 'B']::my_enum[]"), TableReference(table_schema='public', table_name='b', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A'::my_enum, 'B'::my_enum]")], - enum_values_to_rename=[]) - op.drop_column('a', 'value') - # ### end Alembic commands ### + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( # type: ignore[attr-defined] + enum_schema='public', + enum_name='my_enum', + new_values=['A', 'B'], + affected_columns=[TableReference(table_schema='public', table_name='a', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A', 'B']::my_enum[]"), TableReference(table_schema='public', table_name='b', column_name='value', column_type=ColumnType.ARRAY, existing_server_default="ARRAY['A'::my_enum, 'B'::my_enum]")], + enum_values_to_rename=[], + ) + op.drop_column('a', 'value') + # ### end Alembic commands ### """