- Fixed an 0.8 regression whereby the "imports" dictionary member of

the autogen context was removed; this collection is documented in the
"render custom type" documentation as a place to add new imports.
The member is now known as
:attr:`.AutogenContext.imports` and the documentation is repaired.
fixes #332

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 0bdaa970f07f9ac526c59d924a9f30ea6e32641a..61f4036c34b822b0d211fc0e6719ad24c9e73953 100644 (file)
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -220,6 +220,22 @@ class AutogenContext(object):
 
     """
 
+    imports = None
+    """A ``set()`` which contains string Python import directives.
+
+    The directives are to be rendered into the ``${imports}`` section
+    of a script template.  The set is normally empty and can be modified
+    within hooks such as the :paramref:`.EnvironmentContext.configure.render_item`
+    hook.
+
+    .. versionadded:: 0.8.3
+
+    .. seealso::
+
+        :ref:`autogen_render_types`
+
+    """
+
     migration_context = None
     """The :class:`.MigrationContext` established by the ``env.py`` script."""
 
@@ -271,7 +287,7 @@ class AutogenContext(object):
             self.connection = self.migration_context.bind
             self.dialect = self.migration_context.dialect
 
-        self._imports = set()
+        self.imports = set()
         self.opts = opts
         self._has_batch = False
 
@@ -329,9 +345,11 @@ class RevisionContext(object):
         if getattr(migration_script, '_needs_render', False):
             autogen_context = self._last_autogen_context
 
-            autogen_context._imports = set()
+            # clear out existing imports if we are doing multiple
+            # renders
+            autogen_context.imports = set()
             if migration_script.imports:
-                autogen_context._imports.union_update(migration_script.imports)
+                autogen_context.imports.union_update(migration_script.imports)
             render._render_python_into_templatevars(
                 autogen_context, migration_script, template_args
             )

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index e7a616efc8d94d3d1d659c87514a0953d81aec6f..9a64641004af708bfd84a176a50e003e37e4d33c 100644 (file)
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -31,7 +31,7 @@ def _indent(text):
 
 def _render_python_into_templatevars(
         autogen_context, migration_script, template_args):
-    imports = autogen_context._imports
+    imports = autogen_context.imports
 
     for upgrade_ops, downgrade_ops in zip(
             migration_script.upgrade_ops_list,
@@ -567,7 +567,7 @@ def _repr_type(type_, autogen_context):
         return rendered
 
     mod = type(type_).__module__
-    imports = autogen_context._imports
+    imports = autogen_context.imports
     if mod.startswith("sqlalchemy.dialects"):
         dname = re.match(r"sqlalchemy\.dialects\.(\w+)", mod).group(1)
         if imports is not None:

diff --git a/docs/build/autogenerate.rst b/docs/build/autogenerate.rst
index 549aeec4810f9bed4fd3e3310c75fb406db51b05..d0fce8a1bb24f0d5306671ffa4e9932582f8dd30 100644 (file)
--- a/docs/build/autogenerate.rst
+++ b/docs/build/autogenerate.rst
@@ -326,8 +326,8 @@ In the above example, we'd ensure our ``MySpecialType`` includes an appropriate
 ``__repr__()`` method, which is invoked when we call it against ``"%r"``.
 
 The callable we use for :paramref:`.EnvironmentContext.configure.render_item`
-can also add imports to our migration script.  The ``autogen_context`` passed in
-contains an entry called ``autogen_context['imports']``, which is a Python
+can also add imports to our migration script.  The :class:`.AutogenContext` passed in
+contains a datamember called :attr:`.AutogenContext.imports`, which is a Python
 ``set()`` for which we can add new imports.  For example, if ``MySpecialType``
 were in a module called ``mymodel.types``, we can add the import for it
 as we encounter the type::
@@ -337,12 +337,19 @@ as we encounter the type::
 
         if type_ == 'type' and isinstance(obj, MySpecialType):
             # add import for this type
-            autogen_context['imports'].add("from mymodel import types")
+            autogen_context.imports.add("from mymodel import types")
             return "types.%r" % obj
 
         # default rendering for other objects
         return False
 
+.. versionchanged:: 0.8 The ``autogen_context`` data member passed to
+   the ``render_item`` callable is now an instance of :class:`.AutogenContext`.
+
+.. versionchanged:: 0.8.3 The "imports" data member of the autogen context
+   is restored to the new :class:`.AutogenContext` object as
+   :attr:`.AutogenContext.imports`.
+
 The finished migration script will include our imports where the
 ``${imports}`` expression is used, producing output such as::
 

diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index ce917516a93cdf539890cbeb33b23de9080c2dfc..c8e195dc5fee47e665f3a2f7dadfc5b35facf165 100644 (file)
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -6,6 +6,16 @@ Changelog
 .. changelog::
     :version: 0.8.3
 
+    .. change::
+      :tags: bug, autogenerate
+      :tickets: 332
+
+      Fixed an 0.8 regression whereby the "imports" dictionary member of
+      the autogen context was removed; this collection is documented in the
+      "render custom type" documentation as a place to add new imports.
+      The member is now known as
+      :attr:`.AutogenContext.imports` and the documentation is repaired.
+
     .. change::
       :tags: bug, batch
       :tickets: 333

diff --git a/tests/test_autogen_composition.py b/tests/test_autogen_composition.py
index 96d4856e55ce42e08a66b01a6b1552af9926256a..ace6e194bfa22c5052d2652a3507181375ba5d98 100644 (file)
--- a/tests/test_autogen_composition.py
+++ b/tests/test_autogen_composition.py
@@ -205,6 +205,30 @@ nullable=True))
     op.drop_table('item')
     ### end Alembic commands ###""")
 
+    def test_imports_maintined(self):
+        template_args = {}
+        self.context.opts['render_as_batch'] = True
+
+        def render_item(type_, col, autogen_context):
+            autogen_context.imports.add(
+                "from mypackage import my_special_import"
+            )
+            autogen_context.imports.add(
+                "from foobar import bat"
+            )
+
+        self.context.opts["render_item"] = render_item
+        autogenerate._render_migration_diffs(self.context, template_args)
+        eq_(
+            set(
+                template_args['imports'].split("\n")
+            ),
+            set([
+                "from foobar import bat",
+                "from mypackage import my_special_import"
+            ])
+        )
+
 
 class AutogenerateDiffTestWSchema(ModelOne, AutogenTest, TestBase):
     __only_on__ = 'postgresql'

diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index e7010b5a65d128e8be2903ddee4b2bfaafcf028a..c113a233ec724af1fceaa68a5a751bd49092c194 100644 (file)
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -992,14 +992,23 @@ unique=False, """
 
     def test_render_custom(self):
 
+        class MySpecialType(Integer):
+            pass
+
         def render(type_, obj, context):
             if type_ == "foreign_key":
                 return None
             if type_ == "column":
                 if obj.name == "y":
                     return None
+                elif obj.name == "q":
+                    return False
                 else:
                     return "col(%s)" % obj.name
+            if type_ == "type" and isinstance(obj, MySpecialType):
+                context.imports.add("from mypackage import MySpecialType")
+                return "MySpecialType()"
+
             return "render:%s" % type_
 
         self.autogen_context.opts.update(
@@ -1010,6 +1019,7 @@ unique=False, """
         t = Table('t', MetaData(),
                   Column('x', Integer),
                   Column('y', Integer),
+                  Column('q', MySpecialType()),
                   PrimaryKeyConstraint('x'),
                   ForeignKeyConstraint(['x'], ['y'])
                   )
@@ -1019,8 +1029,13 @@ unique=False, """
             result,
             "sa.create_table('t',"
             "col(x),"
+            "sa.Column('q', MySpecialType(), nullable=True),"
             "render:primary_key)"
         )
+        eq_(
+            self.autogen_context.imports,
+            set(['from mypackage import MySpecialType'])
+        )
 
     def test_render_modify_type(self):
         op_obj = ops.AlterColumnOp(
@@ -1350,7 +1365,7 @@ unique=False, """
             autogenerate.render._repr_type(type_, self.autogen_context),
             "mysql.VARCHAR(charset='utf8', national=True, length=20)"
         )
-        eq_(self.autogen_context._imports,
+        eq_(self.autogen_context.imports,
             set(['from sqlalchemy.dialects import mysql'])
             )