From cd0bdca35f9b48d37a1f93c5b036786b01ea0372 Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Fri, 16 Oct 2015 12:29:32 -0400 Subject: [PATCH] - 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 --- alembic/autogenerate/api.py | 24 +++++++++++++++++++++--- alembic/autogenerate/render.py | 4 ++-- docs/build/autogenerate.rst | 13 ++++++++++--- docs/build/changelog.rst | 10 ++++++++++ tests/test_autogen_composition.py | 24 ++++++++++++++++++++++++ tests/test_autogen_render.py | 17 ++++++++++++++++- 6 files changed, 83 insertions(+), 9 deletions(-) diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py index 0bdaa970..61f4036c 100644 --- 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 e7a616ef..9a646410 100644 --- 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 549aeec4..d0fce8a1 100644 --- 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 ce917516..c8e195dc 100644 --- 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 96d4856e..ace6e194 100644 --- 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 e7010b5a..c113a233 100644 --- 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']) ) -- 2.47.2