add user_module_prefix param for render_python_code

Fixed issue where :func:`.autogenerate.render_python_code` function did not
provide a default value for the ``user_module_prefix`` variable, leading to
``NoneType`` errors when autogenerate structures included user-defined
types. Added new parameter
:paramref:`.autogenerate.render_python_code.user_module_prefix` to allow
this to be set as well as to default to ``None``. Pull request courtesy
tangkikodo.

Fixes #1235
Closes: #1233
Pull-request: https://github.com/sqlalchemy/alembic/pull/1233
Pull-request-sha: 6fa3bc1ce4f2130bf7a28282a84e3bb9a3223304

Change-Id: Ic21eec87f373a9518513a6a308f16e5d2b5ebf81

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 95fdedea0aef69929d00dab3f4ce018ab78a154c..c4ec5c1c8b56d206aeda25c4ba89b2675bb7d82d 100644 (file)
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -6,9 +6,10 @@ from typing import Callable
 from typing import Dict
 from typing import Iterator
 from typing import Optional
+from typing import Sequence
 from typing import Set
-from typing import Tuple
 from typing import TYPE_CHECKING
+from typing import Union
 
 from sqlalchemy import inspect
 
@@ -28,6 +29,7 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.schema import SchemaItem
 
     from ..config import Config
+    from ..operations.ops import DowngradeOps
     from ..operations.ops import MigrationScript
     from ..operations.ops import UpgradeOps
     from ..runtime.environment import NameFilterParentNames
@@ -195,13 +197,14 @@ def produce_migrations(
 
 
 def render_python_code(
-    up_or_down_op: UpgradeOps,
+    up_or_down_op: Union[UpgradeOps, DowngradeOps],
     sqlalchemy_module_prefix: str = "sa.",
     alembic_module_prefix: str = "op.",
     render_as_batch: bool = False,
-    imports: Tuple[str, ...] = (),
+    imports: Sequence[str] = (),
     render_item: Optional[RenderItemFn] = None,
     migration_context: Optional[MigrationContext] = None,
+    user_module_prefix: Optional[str] = None,
 ) -> str:
     """Render Python code given an :class:`.UpgradeOps` or
     :class:`.DowngradeOps` object.
@@ -209,12 +212,24 @@ def render_python_code(
     This is a convenience function that can be used to test the
     autogenerate output of a user-defined :class:`.MigrationScript` structure.
 
+    :param up_or_down_op: :class:`.UpgradeOps` or :class:`.DowngradeOps` object
+    :param sqlalchemy_module_prefix: module prefix for SQLAlchemy objects
+    :param alembic_module_prefix: module prefix for Alembic constructs
+    :param render_as_batch: use "batch operations" style for rendering
+    :param imports: sequence of import symbols to add
+    :param render_item: callable to render items
+    :param migration_context: optional :class:`.MigrationContext`
+    :param user_module_prefix: optional string prefix for user-defined types
+
+     .. versionadded:: 1.11.0
+
     """
     opts = {
         "sqlalchemy_module_prefix": sqlalchemy_module_prefix,
         "alembic_module_prefix": alembic_module_prefix,
         "render_item": render_item,
         "render_as_batch": render_as_batch,
+        "user_module_prefix": user_module_prefix,
     }
 
     if migration_context is None:

diff --git a/docs/build/unreleased/1235.rst b/docs/build/unreleased/1235.rst
new file mode 100644 (file)
index 0000000..1a1a887
--- /dev/null
+++ b/docs/build/unreleased/1235.rst
@@ -0,0 +1,11 @@
+.. change::
+    :tags: bug, api, autogenerate
+    :tickets: 1235
+
+    Fixed issue where :func:`.autogenerate.render_python_code` function did not
+    provide a default value for the ``user_module_prefix`` variable, leading to
+    ``NoneType`` errors when autogenerate structures included user-defined
+    types. Added new parameter
+    :paramref:`.autogenerate.render_python_code.user_module_prefix` to allow
+    this to be set as well as to default to ``None``. Pull request courtesy
+    tangkikodo.

diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index 14a0001ca7eac691936cef6e97a1b96190259694..c4f474e84820ddab29ff571a859d822b16fab0ad 100644 (file)
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -1636,6 +1636,34 @@ class AutogenRenderTest(TestBase):
             "sa.ARRAY(sa.DateTime(timezone=True))",
         )
 
+    @config.combinations(None, "default", "my_module.")
+    def test_render_create_table_with_user_module_type(self, mod):
+        class MyType(UserDefinedType):
+            def get_col_spec(self):
+                return "MYTYPE"
+
+        type_ = MyType()
+        uo = ops.UpgradeOps(
+            ops=[ops.CreateTableOp("sometable", [Column("x", type_)])]
+        )
+        if mod != "default":
+            kw = {"user_module_prefix": mod}
+        else:
+            kw = {}
+        if mod and mod != "default":
+            prefix = mod
+        else:
+            prefix = f"{__name__}."
+
+        eq_(
+            autogenerate.render_python_code(uo, **kw),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    op.create_table('sometable',\n"
+            f"    sa.Column('x', {prefix}MyType(), nullable=True)\n"
+            "    )\n"
+            "    # ### end Alembic commands ###",
+        )
+
     def test_render_array_no_context(self):
         uo = ops.UpgradeOps(
             ops=[