Apply subtype repr logic to JSON/JSONB

Fixed bug where Postgresql JSON/JSONB types rendered on SQLAlchemy
1.1 would render the "astext_type" argument which defaults to
the ``Text()`` type without the module prefix, similarly to the
issue with ARRAY fixed in :ticket:`85`.

Also modifies the ARRAY approach from :ticket:`85` to be
regular expression based for safer targeting of the inner
repr() type.

Change-Id: I66d51301f4bf5b747b5e8da26a83cbff075d71b2
Fixes: #411

3 files changed, 71 insertions, 3 deletions

diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py
index 36d9738..9d77ace 100644
--- a/alembic/ddl/postgresql.py
+++ b/alembic/ddl/postgresql.py
@@ -189,11 +189,36 @@ class PostgresqlImpl(DefaultImpl):
 
         return False
 
-    def _render_ARRAY_type(self, type_, autogen_context):
-        sub_type = render._repr_type(type_.item_type, autogen_context)
-        outer_type = repr(type_).replace(repr(type_.item_type), sub_type)
+    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(
+            type_, autogen_context, 'item_type', r'(.+?\()'
+        )
+
+    def _render_JSON_type(self, type_, autogen_context):
+        return self._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(
+            type_, autogen_context, 'astext_type', r'(.+?\(.*astext_type=)'
+        )
+
 
 class PostgresqlColumnType(AlterColumn):
 
diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index 23193ff..5552448 100644
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -24,6 +24,15 @@ Changelog
       flag does **not** support alteration of a column's "autoincrement" status,
       as this is not portable across backends.
 
+    .. change:: 411
+      :tags: bug, postgresql
+      :tickets: 411
+
+      Fixed bug where Postgresql JSON/JSONB types rendered on SQLAlchemy
+      1.1 would render the "astext_type" argument which defaults to
+      the ``Text()`` type without the module prefix, similarly to the
+      issue with ARRAY fixed in :ticket:`85`.
+
     .. change:: 85
       :tags: bug, postgresql
       :tickets: 85
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index 1666e50..e8608b7 100644
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -30,6 +30,10 @@ from sqlalchemy import Boolean
 from sqlalchemy.sql import false
 
 
+if util.sqla_09:
+    from sqlalchemy.dialects.postgresql import JSON, JSONB
+
+
 class PostgresqlOpTest(TestBase):
 
     def test_rename_table_postgresql(self):
@@ -732,3 +736,33 @@ unique=False, """
             "where=sa.text(!U'x != 2'), using='gist', name='t_excl_x')"
             ")"
         )
+
+    @config.requirements.sqlalchemy_09
+    def test_json_type(self):
+        if config.requirements.sqlalchemy_110.enabled:
+            eq_ignore_whitespace(
+                autogenerate.render._repr_type(
+                    JSON(), self.autogen_context),
+                "postgresql.JSON(astext_type=sa.Text())"
+            )
+        else:
+            eq_ignore_whitespace(
+                autogenerate.render._repr_type(
+                    JSON(), self.autogen_context),
+                "postgresql.JSON()"
+            )
+
+    @config.requirements.sqlalchemy_09
+    def test_jsonb_type(self):
+        if config.requirements.sqlalchemy_110.enabled:
+            eq_ignore_whitespace(
+                autogenerate.render._repr_type(
+                    JSONB(), self.autogen_context),
+                "postgresql.JSONB(astext_type=sa.Text())"
+            )
+        else:
+            eq_ignore_whitespace(
+                autogenerate.render._repr_type(
+                    JSONB(), self.autogen_context),
+                "postgresql.JSONB()"
+            )