Add aliases support for get_columns

sqlalchemy_utils/functions/orm.py

@@ -47,7 +47,7 @@ def primary_keys(mixed):
 def get_columns(mixed):
     """
     Return a collection of all Column objects for given SQLAlchemy
-    Table object, declarative class or declarative class instance.
+    object.
 
     The type of the collection depends on the type of the object to return the
     columns from.
@@ -60,12 +60,25 @@ def get_columns(mixed):
 
         get_columns(User.__table__)
 
+        get_columns(User.__mapper__)
+
+        get_column(sa.orm.aliased(User))
+
+        get_columns(sa.orm.alised(User.__table__))
+
 
     :param mixed:
-        SA Table object, SA declarative class or SA declarative class instance
+        SA Table object, SA Mapper, SA declarative class, SA declarative class
+        instance or an alias of any of these objects
     """
     if isinstance(mixed, sa.Table):
         return mixed.c
+    if isinstance(mixed, sa.orm.util.AliasedClass):
+        return sa.inspect(mixed).mapper.columns
+    if isinstance(mixed, sa.sql.selectable.Alias):
+        return mixed.c
+    if isinstance(mixed, sa.orm.Mapper):
+        return mixed.columns
     if not isclass(mixed):
         mixed = mixed.__class__
     return sa.inspect(mixed).columns

tests/functions/test_get_columns.py

@@ -29,3 +29,22 @@ class TestGetColumns(TestCase):
             get_columns(self.Building()),
             sa.util._collections.OrderedProperties
         )
+
+    def test_mapper(self):
+        assert isinstance(
+            get_columns(self.Building.__mapper__),
+            sa.util._collections.OrderedProperties
+        )
+
+    def test_class_alias(self):
+        assert isinstance(
+            get_columns(sa.orm.aliased(self.Building)),
+            sa.util._collections.OrderedProperties
+        )
+
+    def test_table_alias(self):
+        alias = sa.orm.aliased(self.Building.__table__)
+        assert isinstance(
+            get_columns(alias),
+            sa.sql.base.ImmutableColumnCollection
+        )

tests/functions/test_primary_keys.py

@@ -30,3 +30,14 @@ class TestPrimaryKeys(TestCase):
         assert primary_keys(self.Building()) == OrderedDict({
             'id': self.Building.__table__.c._id
         })
+
+    def test_class_alias(self):
+        assert primary_keys(sa.orm.aliased(self.Building())) == OrderedDict({
+            'id': self.Building.__table__.c._id
+        })
+
+    def test_table_alias(self):
+        alias = sa.orm.aliased(self.Building.__table__)
+        assert primary_keys(alias) == OrderedDict({
+            '_id': alias.c._id
+        })