Add ColumnEntity support for get_mapper

CHANGES.rst

@@ -7,7 +7,7 @@ Here you can see the full list of changes between each SQLAlchemy-Utils release.
 0.27.7 (2014-11-03)
 ^^^^^^^^^^^^^^^^^^^
 
-- Added support for Column objects in get_mapper
+- Added support for Column and ColumnEntity objects in get_mapper
 - Made make_order_by_deterministic add deterministic column more aggressively
 
 

sqlalchemy_utils/functions/orm.py

@@ -79,6 +79,8 @@ def get_mapper(mixed):
         mixed = mixed.expr
     elif isinstance(mixed, sa.Column):
         mixed = mixed.table
+    elif isinstance(mixed, sa.orm.query._ColumnEntity):
+        mixed = mixed.expr
 
     if isinstance(mixed, sa.orm.Mapper):
         return mixed

tests/functions/test_get_mapper.py

@@ -91,6 +91,10 @@ class TestGetMapperWithQueryEntities(TestCase):
             sa.inspect(self.Building)
         )
 
+    def test_column_entity(self):
+        query = self.session.query(self.Building.id)
+        assert get_mapper(query._entities[0]) == sa.inspect(self.Building)
+
 
 class TestGetMapperWithMultipleMappersFound(object):
     def setup_method(self, method):