Autogen-Import for Values other than Schema Constructs Dynamically

[pykenny]

I'm trying to make Alembic code auto-generation being able to import customized classes dynamically by following the sample code, but getting a use case that doesn't fit into this example.

Say I'm using mydialect dialect engine, and offering an additional schema option called "customize", where available options are all defined through an enum type:

src/mydialect/enums.py

from enum import Enum

class CustomizedType(Enum):
    def __repr__(self):
        """ Alembic rendering hack """
        return f"CustomizedType.{self.name}"

    A = "A"
    B = "B"
    C = "C"

When I'm defining my ORM model or table, I can pass CustomizedType's value into "customize" argument:

src/models/my_model.py

from sqlalchemy import Column, Integer

from my_orm.base import Base
from my_orm.mydialect.enums import CustomizedType

class MyModel(Base):
    __tablename__ = "my_model"
    __table_args__ = {"mydialect_customize": CustomizedType.C}

    id = Column(Integer, primary_key=True)

And ideally, it would be possible to dynamically add imports when rendering items:

alembic/env.py:

...
from my_orm.mydialect.enums import CustomizedType
...
...
def render_item(type_, item, autogen_context):
    if isinstance(item, CustomizedType):
        autogen_context.imports.add("from my_orm.mydialect.enums import CustomizedType")

    return False
...

and finally get some migration script that looks like:

...
from my_orm.mydialect.enums import CustomizedType
...

def upgrade():
    op.create_table(
        "my_model",
        ...
        mydialect_customize=CustomizedType.C,
    )

But from experiments, it seemed that table arguments are never passed to render_item() function. Any suggestions for dealing with importing types/classes that doesn't get passed into the rendering function?

[zzzeek]

that documentation example you point to refers to the type objects you'd have inside of the Column, like an Integer or String. The render_item() hook is only set up for types, columns and constraints right now, not for a CreateTable nor would it be called for individual dialect items inside of the Table's kwargs.

The elements that are inside the Table's "kw" collection are just rendered through repr() right now so you wouldn't get any access to the table or anything like that at this level.

to get a top level view of the whole autogenerate, where you can search for things like CreateTableOp and then add imports to the migration context would be the process_revision_directives hook. Docs for this are at https://alembic.sqlalchemy.org/en/latest/api/autogenerate.html#customizing-revision-generation .

that said, imports are cheap I'd likely just have them be static in your script.py.mako template.

[pykenny]

Oh okay, then maybe I'll hard-code all possible imports into the template, then let zimports clean up the unused ones.