Add generic per-type rendering, implement ARRAY

Fixed bug expanding upon the fix for
:ticket:`85` which adds the correct module import to the
"inner" type for an ``ARRAY`` type, the fix now accommodates for the
generic ``sqlalchemy.types.ARRAY`` type added in SQLAlchemy 1.1,
rendering the inner type correctly regardless of whether or not the
Postgresql dialect is present.

Change-Id: I98ffbf95a88dc815404a2d1020b0e3e56742ca8c
Fixes: #442

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 6f5ffed110750ccd6ed31de273a47564eba17d16..d6f398f54223ad087b27c79ec01baf99ecb935d4 100644 (file)
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -573,6 +573,8 @@ def _repr_type(type_, autogen_context):
     if hasattr(autogen_context.migration_context, 'impl'):
         impl_rt = autogen_context.migration_context.impl.render_type(
             type_, autogen_context)
+    else:
+        impl_rt = None
 
     mod = type(type_).__module__
     imports = autogen_context.imports
@@ -585,13 +587,47 @@ def _repr_type(type_, autogen_context):
         else:
             return "%s.%r" % (dname, type_)
     elif mod.startswith("sqlalchemy."):
-        prefix = _sqlalchemy_autogenerate_prefix(autogen_context)
-        return "%s%r" % (prefix, type_)
+        if '_render_%s_type' % type_.__visit_name__ in globals():
+            fn = globals()['_render_%s_type' % type_.__visit_name__]
+            return fn(type_, autogen_context)
+        else:
+            prefix = _sqlalchemy_autogenerate_prefix(autogen_context)
+            return "%s%r" % (prefix, type_)
     else:
         prefix = _user_autogenerate_prefix(autogen_context, type_)
         return "%s%r" % (prefix, type_)
 
 
+def _render_ARRAY_type(type_, autogen_context):
+    return _render_type_w_subtype(
+        type_, autogen_context, 'item_type', r'(.+?\()'
+    )
+
+
+def _render_type_w_subtype(type_, autogen_context, attrname, regexp):
+    outer_repr = repr(type_)
+    inner_type = getattr(type_, attrname, None)
+    if inner_type is None:
+        return False
+
+    inner_repr = repr(inner_type)
+
+    inner_repr = re.sub(r'([\(\)])', r'\\\1', inner_repr)
+    sub_type = _repr_type(getattr(type_, attrname), autogen_context)
+    outer_type = re.sub(
+        regexp + inner_repr,
+        r"\1%s" % sub_type, outer_repr)
+
+    mod = type(type_).__module__
+    if mod.startswith("sqlalchemy.dialects"):
+        dname = re.match(r"sqlalchemy\.dialects\.(\w+)", mod).group(1)
+        return "%s.%s" % (dname, outer_type)
+    elif mod.startswith("sqlalchemy"):
+        prefix = _sqlalchemy_autogenerate_prefix(autogen_context)
+        return "%s%s" % (prefix, outer_type)
+    else:
+        return None
+
 _constraint_renderers = util.Dispatcher()
 
 

diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py
index 9d77acecbcf48d20de2ef6e3dfbcce73d4104bdf..83471aa31318e2983f8b79f6db82f87ecabe3ce4 100644 (file)
--- a/alembic/ddl/postgresql.py
+++ b/alembic/ddl/postgresql.py
@@ -16,7 +16,6 @@ from ..operations.base import BatchOperations
 from ..operations import ops
 from ..util import sqla_compat
 from ..operations import schemaobj
-from ..autogenerate import render
 
 import logging
 
@@ -183,39 +182,28 @@ class PostgresqlImpl(DefaultImpl):
                     metadata_indexes.discard(idx)
 
     def render_type(self, type_, autogen_context):
+        mod = type(type_).__module__
+        if not mod.startswith("sqlalchemy.dialects.postgresql"):
+            return False
+
         if hasattr(self, '_render_%s_type' % type_.__visit_name__):
             meth = getattr(self, '_render_%s_type' % type_.__visit_name__)
             return meth(type_, autogen_context)
 
         return False
 
-    def _render_type_w_subtype(self, type_, autogen_context, attrname, regexp):
-        outer_repr = repr(type_)
-        inner_type = getattr(type_, attrname, None)
-        if inner_type is None:
-            return False
-
-        inner_repr = repr(inner_type)
-
-        inner_repr = re.sub(r'([\(\)])', r'\\\1', inner_repr)
-        sub_type = render._repr_type(getattr(type_, attrname), autogen_context)
-        outer_type = re.sub(
-            regexp + inner_repr,
-            r"\1%s" % sub_type, outer_repr)
-        return "%s.%s" % ("postgresql", outer_type)
-
     def _render_ARRAY_type(self, type_, autogen_context):
-        return self._render_type_w_subtype(
+        return render._render_type_w_subtype(
             type_, autogen_context, 'item_type', r'(.+?\()'
         )
 
     def _render_JSON_type(self, type_, autogen_context):
-        return self._render_type_w_subtype(
+        return render._render_type_w_subtype(
             type_, autogen_context, 'astext_type', r'(.+?\(.*astext_type=)'
         )
 
     def _render_JSONB_type(self, type_, autogen_context):
-        return self._render_type_w_subtype(
+        return render._render_type_w_subtype(
             type_, autogen_context, 'astext_type', r'(.+?\(.*astext_type=)'
         )
 

diff --git a/docs/build/unreleased/442.rst b/docs/build/unreleased/442.rst
new file mode 100644 (file)
index 0000000..0ac7550
--- /dev/null
+++ b/docs/build/unreleased/442.rst
@@ -0,0 +1,10 @@
+.. change::
+    :tags: bug, autogenerate
+    :tickets: 442
+
+    Fixed bug expanding upon the fix for
+    :ticket:`85` which adds the correct module import to the
+    "inner" type for an ``ARRAY`` type, the fix now accommodates for the
+    generic ``sqlalchemy.types.ARRAY`` type added in SQLAlchemy 1.1,
+    rendering the inner type correctly regardless of whether or not the
+    Postgresql dialect is present.
\ No newline at end of file

diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index 2882858a9f677b9349cb4b5474a5c6bc525c84d4..dce7b7157ca5142b58da79a8b76cac6e3a6cd238 100644 (file)
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -7,10 +7,11 @@ from alembic.operations import ops
 from sqlalchemy import MetaData, Column, Table, String, \
     Numeric, CHAR, ForeignKey, DATETIME, Integer, BigInteger, \
     CheckConstraint, Unicode, Enum, cast,\
-    UniqueConstraint, Boolean, ForeignKeyConstraint,\
+    DateTime, UniqueConstraint, Boolean, ForeignKeyConstraint,\
     PrimaryKeyConstraint, Index, func, text, DefaultClause
 
 from sqlalchemy.types import TIMESTAMP
+from sqlalchemy import types
 from sqlalchemy.types import UserDefinedType
 from sqlalchemy.engine.default import DefaultDialect
 from sqlalchemy.sql import and_, column, literal_column, false, table
@@ -1323,6 +1324,39 @@ class AutogenRenderTest(TestBase):
             "sa.Integer()"
         )
 
+    @config.requirements.sqlalchemy_110
+    def test_generic_array_type(self):
+
+        eq_ignore_whitespace(
+            autogenerate.render._repr_type(
+                types.ARRAY(Integer), self.autogen_context),
+            "sa.ARRAY(sa.Integer())"
+        )
+
+        eq_ignore_whitespace(
+            autogenerate.render._repr_type(
+                types.ARRAY(DateTime(timezone=True)), self.autogen_context),
+            "sa.ARRAY(sa.DateTime(timezone=True))"
+        )
+
+    @config.requirements.sqlalchemy_110
+    def test_render_array_no_context(self):
+        uo = ops.UpgradeOps(ops=[
+            ops.CreateTableOp(
+                "sometable",
+                [Column('x', types.ARRAY(Integer))]
+            )
+        ])
+
+        eq_(
+            autogenerate.render_python_code(uo),
+            "# ### commands auto generated by Alembic - please adjust! ###\n"
+            "    op.create_table('sometable',\n"
+            "    sa.Column('x', sa.ARRAY(sa.Integer()), nullable=True)\n"
+            "    )\n"
+            "    # ### end Alembic commands ###"
+        )
+
     def test_repr_custom_type_w_sqla_prefix(self):
         self.autogen_context.opts['user_module_prefix'] = None
 

diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index afa9e7c26a36706a5faf1ab2a4a068a4d2c6c7a2..fd0a182663c3a06ee7a2de470377bdbea4f648d1 100644 (file)
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -3,6 +3,7 @@ from sqlalchemy import DateTime, MetaData, Table, Column, text, Integer, \
     String, Interval, Sequence, Numeric, BigInteger, Float, Numeric
 from sqlalchemy.dialects.postgresql import ARRAY, UUID, BYTEA
 from sqlalchemy.engine.reflection import Inspector
+from sqlalchemy import types
 from alembic.operations import Operations
 from sqlalchemy.sql import table, column
 from alembic.autogenerate.compare import \
@@ -664,7 +665,7 @@ unique=False, """
         )
 
     @config.requirements.sqlalchemy_09
-    def test_array_type(self):
+    def test_postgresql_array_type(self):
 
         eq_ignore_whitespace(
             autogenerate.render._repr_type(
@@ -688,6 +689,34 @@ unique=False, """
         assert 'from sqlalchemy.dialects import postgresql' in \
             self.autogen_context.imports
 
+    @config.requirements.sqlalchemy_110
+    def test_generic_array_type(self):
+
+        eq_ignore_whitespace(
+            autogenerate.render._repr_type(
+                types.ARRAY(Integer), self.autogen_context),
+            "sa.ARRAY(sa.Integer())"
+        )
+
+        eq_ignore_whitespace(
+            autogenerate.render._repr_type(
+                types.ARRAY(DateTime(timezone=True)), self.autogen_context),
+            "sa.ARRAY(sa.DateTime(timezone=True))"
+        )
+
+        assert 'from sqlalchemy.dialects import postgresql' not in \
+            self.autogen_context.imports
+
+        eq_ignore_whitespace(
+            autogenerate.render._repr_type(
+                types.ARRAY(BYTEA, as_tuple=True, dimensions=2),
+                self.autogen_context),
+            "sa.ARRAY(postgresql.BYTEA(), as_tuple=True, dimensions=2)"
+        )
+
+        assert 'from sqlalchemy.dialects import postgresql' in \
+            self.autogen_context.imports
+
     @config.requirements.sqlalchemy_09
     def test_array_type_user_defined_inner(self):
         def repr_type(typestring, object_, autogen_context):