Fixed array support for SQLAlchemy < 1.1

agronholm · agronholm · commit 7520627c8ed8 · 2018-05-20T17:36:43.000+03:00
diff --git a/sqlacodegen/codegen.py b/sqlacodegen/codegen.py
@@ -11,12 +11,18 @@
 import sqlalchemy
 from sqlalchemy import (
     Enum, ForeignKeyConstraint, PrimaryKeyConstraint, CheckConstraint, UniqueConstraint, Table,
-    Column, ARRAY)
+    Column)
 from sqlalchemy.schema import ForeignKey
 from sqlalchemy.sql.expression import TextClause
 from sqlalchemy.types import Boolean, String
 from sqlalchemy.util import OrderedDict
 
+# The generic ARRAY type was introduced in SQLAlchemy 1.1
+try:
+    from sqlalchemy import ARRAY
+except ImportError:
+    from sqlalchemy.dialects.postgresql import ARRAY
+
 # Conditionally import Geoalchemy2 to enable reflection support
 try:
     import geoalchemy2  # noqa: F401
diff --git a/tests/test_codegen.py b/tests/test_codegen.py
@@ -4,6 +4,7 @@
 import sys
 from io import StringIO
 
+import sqlalchemy
 from sqlalchemy.dialects.mysql import base as mysql
 from sqlalchemy.dialects.mysql.base import TINYINT
 from sqlalchemy.dialects.postgresql import ARRAY
@@ -93,7 +94,23 @@ def test_arrays(self):
             Column('int_array', ARRAY(INTEGER))
         )
 
-        assert self.generate_code() == """\
+        if sqlalchemy.__version__ < '1.1':
+            assert self.generate_code() == """\
+# coding: utf-8
+from sqlalchemy import Column, Float, Integer, MetaData, Table
+from sqlalchemy.dialects.postgresql.base import ARRAY
+
+metadata = MetaData()
+
+
+t_simple_items = Table(
+    'simple_items', metadata,
+    Column('dp_array', ARRAY(Float(precision=53))),
+    Column('int_array', ARRAY(Integer()))
+)
+"""
+        else:
+            assert self.generate_code() == """\
 # coding: utf-8
 from sqlalchemy import ARRAY, Column, Float, Integer, MetaData, Table
 
