Skip to content

Commit

Permalink
Merge pull request #56 from Pogchamp-company/feature/-/unify-formatting
Browse files Browse the repository at this point in the history
Unify formatting with black
  • Loading branch information
RustyGuard authored Jan 6, 2024
2 parents a89aab7 + f39dd28 commit ef2d3d8
Show file tree
Hide file tree
Showing 42 changed files with 1,109 additions and 705 deletions.
10 changes: 10 additions & 0 deletions .github/workflows/black.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
name: Lint

on: [push, pull_request]

jobs:
lint:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
- uses: psf/black@stable
22 changes: 14 additions & 8 deletions alembic_postgresql_enum/add_create_type_false.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,14 @@
import logging

import sqlalchemy
from alembic.operations.ops import UpgradeOps, ModifyTableOps, AddColumnOp, CreateTableOp, DropColumnOp, DropTableOp
from alembic.operations.ops import (
UpgradeOps,
ModifyTableOps,
AddColumnOp,
CreateTableOp,
DropColumnOp,
DropTableOp,
)
from sqlalchemy import Column
from sqlalchemy.dialects import postgresql

Expand All @@ -10,13 +17,12 @@ class ReprWorkaround(postgresql.ENUM):
"""
As postgresql.ENUM does not include create_type inside __repr__, we have to swap it with custom type
"""
__module__ = 'sqlalchemy.dialects.postgresql'

__module__ = "sqlalchemy.dialects.postgresql"

def __repr__(self):
return (
f'{super().__repr__()[:-1]}, create_type=False)'
.replace('ReprWorkaround', 'ENUM')
.replace(', metadata=MetaData()', '')
return f"{super().__repr__()[:-1]}, create_type=False)".replace("ReprWorkaround", "ENUM").replace(
", metadata=MetaData()", ""
)


Expand All @@ -26,7 +32,7 @@ def inject_repr_into_enums(column: Column):
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'))
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)
Expand All @@ -36,7 +42,7 @@ def inject_repr_into_enums(column: Column):
column.type = replacement_enum_type


log = logging.getLogger(f'alembic.{__name__}')
log = logging.getLogger(f"alembic.{__name__}")


def add_create_type_false(upgrade_ops: UpgradeOps):
Expand Down
12 changes: 5 additions & 7 deletions alembic_postgresql_enum/add_postgres_using_to_text.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,15 @@ class PostgresUsingAlterColumnOp(AlterColumnOp):

def reverse(self):
reversed_object = super().reverse()
reversed_object.kw.pop('postgresql_using', None)
reversed_object.kw.pop("postgresql_using", None)
return reversed_object


@renderers.dispatch_for(PostgresUsingAlterColumnOp)
def _postgres_using_alter_column(
autogen_context: AutogenContext, op: ops.AlterColumnOp
) -> str:
def _postgres_using_alter_column(autogen_context: AutogenContext, op: ops.AlterColumnOp) -> str:
alter_column_expression = render._alter_column(autogen_context, op)

postgresql_using = op.kw.get('postgresql_using', None)
postgresql_using = op.kw.get("postgresql_using", None)
indent = " " * 11

# To remove closing bracket
Expand All @@ -37,11 +35,11 @@ def _postgres_using_alter_column(
return alter_column_expression


log = logging.getLogger(f'alembic.{__name__}')
log = logging.getLogger(f"alembic.{__name__}")


def add_postgres_using_to_alter_operation(op: AlterColumnOp):
op.kw['postgresql_using'] = f'{op.column_name}::{op.modify_type.name}'
op.kw["postgresql_using"] = f"{op.column_name}::{op.modify_type.name}"
log.info("postgresql_using added to %r.%r alteration", op.table_name, op.column_name)
op.__class__ = PostgresUsingAlterColumnOp

Expand Down
25 changes: 20 additions & 5 deletions alembic_postgresql_enum/compare_dispatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,23 @@
from alembic.operations.ops import UpgradeOps, CreateTableOp

from alembic_postgresql_enum.add_create_type_false import add_create_type_false
from alembic_postgresql_enum.add_postgres_using_to_text import add_postgres_using_to_text
from alembic_postgresql_enum.detection_of_changes import sync_changed_enums, create_new_enums, drop_unused_enums
from alembic_postgresql_enum.add_postgres_using_to_text import (
add_postgres_using_to_text,
)
from alembic_postgresql_enum.detection_of_changes import (
sync_changed_enums,
create_new_enums,
drop_unused_enums,
)
from alembic_postgresql_enum.get_enum_data import get_defined_enums, get_declared_enums


@alembic.autogenerate.comparators.dispatch_for("schema")
def compare_enums(autogen_context: AutogenContext, upgrade_ops: UpgradeOps, schema_names: Iterable[Union[str, None]]):
def compare_enums(
autogen_context: AutogenContext,
upgrade_ops: UpgradeOps,
schema_names: Iterable[Union[str, None]],
):
"""
Walk the declared SQLAlchemy schema for every referenced Enum, walk the PG
schema for every defined Enum, then generate SyncEnumValuesOp migrations
Expand Down Expand Up @@ -41,5 +51,10 @@ def compare_enums(autogen_context: AutogenContext, upgrade_ops: UpgradeOps, sche

drop_unused_enums(definitions, declarations.enum_values, schema, upgrade_ops)

sync_changed_enums(definitions, declarations.enum_values,
declarations.enum_table_references, schema, upgrade_ops)
sync_changed_enums(
definitions,
declarations.enum_values,
declarations.enum_table_references,
schema,
upgrade_ops,
)
36 changes: 25 additions & 11 deletions alembic_postgresql_enum/detection_of_changes/enum_alteration.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,22 @@

from alembic.operations.ops import UpgradeOps

from alembic_postgresql_enum.get_enum_data import EnumNamesToValues, EnumNamesToTableReferences
from alembic_postgresql_enum.get_enum_data import (
EnumNamesToValues,
EnumNamesToTableReferences,
)
from alembic_postgresql_enum.operations.sync_enum_values import SyncEnumValuesOp

log = logging.getLogger(f'alembic.{__name__}')
log = logging.getLogger(f"alembic.{__name__}")


def sync_changed_enums(defined_enums: EnumNamesToValues,
declared_enums: EnumNamesToValues,
table_references: EnumNamesToTableReferences,
schema: str,
upgrade_ops: UpgradeOps,
):
def sync_changed_enums(
defined_enums: EnumNamesToValues,
declared_enums: EnumNamesToValues,
table_references: EnumNamesToTableReferences,
schema: str,
upgrade_ops: UpgradeOps,
):
for enum_name, new_values in declared_enums.items():
if enum_name not in defined_enums:
# That is work for create_new_enums function
Expand All @@ -30,8 +34,18 @@ def sync_changed_enums(defined_enums: EnumNamesToValues,
# Enum definition and declaration are in sync
continue

log.info("Detected changed enum values in %r\nWas: %r\nBecome: %r", enum_name,
list(old_values), list(new_values))
log.info(
"Detected changed enum values in %r\nWas: %r\nBecome: %r",
enum_name,
list(old_values),
list(new_values),
)
affected_columns = table_references[enum_name]
op = SyncEnumValuesOp(schema, enum_name, list(old_values), list(new_values), list(affected_columns))
op = SyncEnumValuesOp(
schema,
enum_name,
list(old_values),
list(new_values),
list(affected_columns),
)
upgrade_ops.ops.append(op)
10 changes: 7 additions & 3 deletions alembic_postgresql_enum/detection_of_changes/enum_creation.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,15 @@
from alembic_postgresql_enum.get_enum_data import EnumNamesToValues
from alembic_postgresql_enum.operations.create_enum import CreateEnumOp

log = logging.getLogger(f'alembic.{__name__}')
log = logging.getLogger(f"alembic.{__name__}")


def create_new_enums(defined_enums: EnumNamesToValues, declared_enums: EnumNamesToValues,
schema: str, upgrade_ops: UpgradeOps):
def create_new_enums(
defined_enums: EnumNamesToValues,
declared_enums: EnumNamesToValues,
schema: str,
upgrade_ops: UpgradeOps,
):
"""
Create enums that are not in Postgres schema
"""
Expand Down
10 changes: 7 additions & 3 deletions alembic_postgresql_enum/detection_of_changes/enum_deletion.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,15 @@
from alembic_postgresql_enum.get_enum_data import EnumNamesToValues
from alembic_postgresql_enum.operations.drop_enum import DropEnumOp

log = logging.getLogger(f'alembic.{__name__}')
log = logging.getLogger(f"alembic.{__name__}")


def drop_unused_enums(defined_enums: EnumNamesToValues, declared_enums: EnumNamesToValues,
schema: str, upgrade_ops: UpgradeOps):
def drop_unused_enums(
defined_enums: EnumNamesToValues,
declared_enums: EnumNamesToValues,
schema: str,
upgrade_ops: UpgradeOps,
):
"""
Drop enums that are in Postgres schema but not declared in SqlAlchemy schema
"""
Expand Down
8 changes: 7 additions & 1 deletion alembic_postgresql_enum/get_enum_data/__init__.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,9 @@
from .types import ColumnType, DeclaredEnumValues, EnumNamesToValues, EnumNamesToTableReferences, TableReference
from .types import (
ColumnType,
DeclaredEnumValues,
EnumNamesToValues,
EnumNamesToTableReferences,
TableReference,
)
from .defined_enums import get_defined_enums
from .declared_enums import get_declared_enums
36 changes: 21 additions & 15 deletions alembic_postgresql_enum/get_enum_data/declared_enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,22 +10,26 @@
if TYPE_CHECKING:
from sqlalchemy.engine import Connection

from alembic_postgresql_enum.get_enum_data import DeclaredEnumValues, TableReference, ColumnType
from alembic_postgresql_enum.get_enum_data import (
DeclaredEnumValues,
TableReference,
ColumnType,
)


def get_enum_values(enum_type: sqlalchemy.Enum) -> 'Tuple[str, ...]':
def get_enum_values(enum_type: sqlalchemy.Enum) -> "Tuple[str, ...]":
# For specific case when types.TypeDecorator is used
if isinstance(enum_type, sqlalchemy.types.TypeDecorator):
dialect = postgresql.dialect

def value_processor(value):
return enum_type.process_bind_param(
enum_type.impl.result_processor(dialect, enum_type)(value),
dialect
)
return enum_type.process_bind_param(enum_type.impl.result_processor(dialect, enum_type)(value), dialect)

else:

def value_processor(enum_value):
return enum_value

return tuple(value_processor(value) for value in enum_type.enums)


Expand All @@ -34,17 +38,18 @@ def column_type_is_enum(column_type: Any) -> bool:
return column_type.native_enum

# For specific case when types.TypeDecorator is used
if isinstance(getattr(column_type, 'impl', None), sqlalchemy.Enum):
if isinstance(getattr(column_type, "impl", None), sqlalchemy.Enum):
return True

return False


def get_declared_enums(metadata: Union[MetaData, List[MetaData]],
schema: str,
default_schema: str,
connection: 'Connection',
) -> DeclaredEnumValues:
def get_declared_enums(
metadata: Union[MetaData, List[MetaData]],
schema: str,
default_schema: str,
connection: "Connection",
) -> DeclaredEnumValues:
"""
Return a dict mapping SQLAlchemy declared enumeration types to the set of their values
with columns where enums are used.
Expand Down Expand Up @@ -102,7 +107,8 @@ def get_declared_enums(metadata: Union[MetaData, List[MetaData]],

return DeclaredEnumValues(
enum_values=enum_name_to_values,
enum_table_references={enum_name: frozenset(table_references)
for enum_name, table_references
in enum_name_to_table_references.items()},
enum_table_references={
enum_name: frozenset(table_references)
for enum_name, table_references in enum_name_to_table_references.items()
},
)
11 changes: 4 additions & 7 deletions alembic_postgresql_enum/get_enum_data/defined_enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,15 @@


def _remove_schema_prefix(enum_name: str, schema: str) -> str:
schema_prefix = f'{schema}.'
schema_prefix = f"{schema}."

if enum_name.startswith(schema_prefix):
enum_name = enum_name[len(schema_prefix):]
enum_name = enum_name[len(schema_prefix) :]

return enum_name


def get_defined_enums(connection: 'Connection', schema: str) -> EnumNamesToValues:
def get_defined_enums(connection: "Connection", schema: str) -> EnumNamesToValues:
"""
Return a dict mapping PostgreSQL defined enumeration types to the set of their
defined values.
Expand All @@ -29,7 +29,4 @@ def get_defined_enums(connection: 'Connection', schema: str) -> EnumNamesToValue
"my_enum": tuple(["a", "b", "c"]),
}
"""
return {
_remove_schema_prefix(name, schema): tuple(values)
for name, values in get_all_enums(connection, schema)
}
return {_remove_schema_prefix(name, schema): tuple(values) for name, values in get_all_enums(connection, schema)}
14 changes: 7 additions & 7 deletions alembic_postgresql_enum/get_enum_data/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ class ColumnType(PyEnum):
ARRAY = ARRAY

def __repr__(self):
return f'{self.__class__.__name__}.{self.name}'
return f"{self.__class__.__name__}.{self.name}"


@dataclass(frozen=True)
Expand All @@ -21,15 +21,15 @@ class TableReference:
existing_server_default: str = None

def __repr__(self):
result_str = 'TableReference('
result_str += f'{self.table_name!r}, '
result_str += f'{self.column_name!r}, '
result_str = "TableReference("
result_str += f"{self.table_name!r}, "
result_str += f"{self.column_name!r}, "
if self.column_type != ColumnType.COMMON:
result_str += f'column_type=ColumnType.{self.column_type.name}, '
result_str += f"column_type=ColumnType.{self.column_type.name}, "
if self.existing_server_default is not None:
result_str += f'existing_server_default={self.existing_server_default!r}, '
result_str += f"existing_server_default={self.existing_server_default!r}, "
result_str = result_str[:-2] # to remove last comma with space
result_str += ')'
result_str += ")"
return result_str

@property
Expand Down
3 changes: 2 additions & 1 deletion alembic_postgresql_enum/operations/create_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@


class CreateEnumOp(EnumLifecycleOp):
operation_name = 'create_enum'
operation_name = "create_enum"

def reverse(self):
from .drop_enum import DropEnumOp

return DropEnumOp(
name=self.name,
schema=self.schema,
Expand Down
Loading

0 comments on commit ef2d3d8

Please sign in to comment.