From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Thu, 19 Sep 2019 21:40:06 +0000 (-0400)
Subject: Render Variant types
X-Git-Tag: rel_1_2_0~2^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f8c1c70eb9cea8c7fb7ef1236e29b924908c5b8b;p=thirdparty%2Fsqlalchemy%2Falembic.git

Render Variant types

Added rendering for SQLAlchemy ``Variant`` datatypes, which render as the
base type plus one or more ``.with_variant()`` method calls.

Change-Id: Ic0eea1be911c2fb3e9802ec21ff99815a06602e3
Fixes: #131
---

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 895c3903..81016370 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -679,6 +679,8 @@ def _repr_type(type_, autogen_context):
     elif impl_rt:
         return impl_rt
     elif mod.startswith("sqlalchemy."):
+        if type(type_) is sqltypes.Variant:
+            return _render_Variant_type(type_, autogen_context)
         if "_render_%s_type" % type_.__visit_name__ in globals():
             fn = globals()["_render_%s_type" % type_.__visit_name__]
             return fn(type_, autogen_context)
@@ -696,6 +698,17 @@ def _render_ARRAY_type(type_, autogen_context):
     )
 
 
+def _render_Variant_type(type_, autogen_context):
+    base = _repr_type(type_.impl, autogen_context)
+    for dialect in sorted(type_.mapping):
+        typ = type_.mapping[dialect]
+        base += ".with_variant(%r, %s)" % (
+            dialect,
+            _repr_type(typ, autogen_context),
+        )
+    return base
+
+
 def _render_type_w_subtype(
     type_, autogen_context, attrname, regexp, prefix=None
 ):
diff --git a/docs/build/unreleased/131.rst b/docs/build/unreleased/131.rst
new file mode 100644
index 00000000..f37edcec
--- /dev/null
+++ b/docs/build/unreleased/131.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: usecase, autogenerate
+    :tickets: 131
+
+    Added rendering for SQLAlchemy ``Variant`` datatypes, which render as the
+    base type plus one or more ``.with_variant()`` method calls.
+
diff --git a/tests/test_autogen_render.py b/tests/test_autogen_render.py
index b31c3f8f..6202a4f2 100644
--- a/tests/test_autogen_render.py
+++ b/tests/test_autogen_render.py
@@ -1677,6 +1677,25 @@ class AutogenRenderTest(TestBase):
             "sqlalchemy_util.types.MyType()",
         )
 
+    def test_render_variant(self):
+        from sqlalchemy import VARCHAR, CHAR
+
+        self.autogen_context.opts["user_module_prefix"] = None
+
+        type_ = (
+            String(5)
+            .with_variant(VARCHAR(10), "mysql")
+            .with_variant(CHAR(15), "oracle")
+        )
+
+        # the new Black formatting will help a lot with this
+        eq_ignore_whitespace(
+            autogenerate.render._repr_type(type_, self.autogen_context),
+            "sa.String(length=5)."
+            "with_variant('mysql', sa.VARCHAR(length=10))."
+            "with_variant('oracle', sa.CHAR(length=15))",
+        )
+
     def test_repr_user_type_user_prefix_None(self):
         class MyType(UserDefinedType):
             def get_col_spec(self):