support 2.0 style variants

Implemented support for recognizing and rendering SQLAlchemy "variant"
types going forward into SQLAlchemy 2.0, where the architecture of
"variant" datatypes will be changing.

Specifically any TypeEngine can now have "variants" by looking
in the _variant_mapping dictionary collection associated with the
type directly.

Change-Id: I0b5ca887dce1dd77af3504dbe318701b31c7d574

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 90d49e5f373002a620b836618ebde12bb20d6e43..b8226f7a7007bcbf00fc59bb3adcda419af69847 100644 (file)
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -41,7 +41,6 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.schema import UniqueConstraint
     from sqlalchemy.sql.sqltypes import ARRAY
     from sqlalchemy.sql.type_api import TypeEngine
-    from sqlalchemy.sql.type_api import Variant
 
     from alembic.autogenerate.api import AutogenContext
     from alembic.config import Config
@@ -812,7 +811,11 @@ def _get_identity_options(identity_options: "Identity") -> OrderedDict:
     return kwargs
 
 
-def _repr_type(type_: "TypeEngine", autogen_context: "AutogenContext") -> str:
+def _repr_type(
+    type_: "TypeEngine",
+    autogen_context: "AutogenContext",
+    _skip_variants: bool = False,
+) -> str:
     rendered = _user_defined_render("type", type_, autogen_context)
     if rendered is not False:
         return rendered
@@ -839,7 +842,7 @@ def _repr_type(type_: "TypeEngine", autogen_context: "AutogenContext") -> str:
     elif impl_rt:
         return impl_rt
     elif mod.startswith("sqlalchemy."):
-        if type(type_) is sqltypes.Variant:
+        if not _skip_variants and sqla_compat._type_has_variants(type_):
             return _render_Variant_type(type_, autogen_context)
         if "_render_%s_type" % type_.__visit_name__ in globals():
             fn = globals()["_render_%s_type" % type_.__visit_name__]
@@ -864,14 +867,15 @@ def _render_ARRAY_type(
 
 
 def _render_Variant_type(
-    type_: "Variant", autogen_context: "AutogenContext"
+    type_: "TypeEngine", autogen_context: "AutogenContext"
 ) -> str:
-    base = _repr_type(type_.impl, autogen_context)
+    base_type, variant_mapping = sqla_compat._get_variant_mapping(type_)
+    base = _repr_type(base_type, autogen_context, _skip_variants=True)
     assert base is not None and base is not False
-    for dialect in sorted(type_.mapping):
-        typ = type_.mapping[dialect]
+    for dialect in sorted(variant_mapping):
+        typ = variant_mapping[dialect]
         base += ".with_variant(%s, %r)" % (
-            _repr_type(typ, autogen_context),
+            _repr_type(typ, autogen_context, _skip_variants=True),
             dialect,
         )
     return base

diff --git a/alembic/util/sqla_compat.py b/alembic/util/sqla_compat.py
index 57a6a769dcbcbbfe0ff4f4e86008122850f5261e..221e20e86c71f7b91114d7c44526e0f49719fa1c 100644 (file)
--- a/alembic/util/sqla_compat.py
+++ b/alembic/util/sqla_compat.py
@@ -252,6 +252,24 @@ def _reflect_table(
         return inspector.reflecttable(table, None)
 
 
+if hasattr(sqltypes.TypeEngine, "_variant_mapping"):
+
+    def _type_has_variants(type_):
+        return bool(type_._variant_mapping)
+
+    def _get_variant_mapping(type_):
+        return type_, type_._variant_mapping
+
+
+else:
+
+    def _type_has_variants(type_):
+        return type(type_) is sqltypes.Variant
+
+    def _get_variant_mapping(type_):
+        return type_.impl, type_.mapping
+
+
 def _fk_spec(constraint):
     source_columns = [
         constraint.columns[key].name for key in constraint.column_keys

diff --git a/docs/build/unreleased/variant.rst b/docs/build/unreleased/variant.rst
new file mode 100644 (file)
index 0000000..71cf3c7
--- /dev/null
+++ b/docs/build/unreleased/variant.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: bug, autogenerate
+
+    Implemented support for recognizing and rendering SQLAlchemy "variant"
+    types going forward into SQLAlchemy 2.0, where the architecture of
+    "variant" datatypes will be changing.
+