Add get_columns utility function

CHANGES.rst

@@ -4,6 +4,13 @@ Changelog
 Here you can see the full list of changes between each SQLAlchemy-Utils release.
 
 
+0.25.3 (2014-04-21)
+^^^^^^^^^^^^^^^^^^^
+
+- Added support for primary key aliases in primary_keys function
+- Added get_columns utility function
+
+
 0.25.2 (2014-03-25)
 ^^^^^^^^^^^^^^^^^^^
 

docs/model_helpers.rst

@@ -16,6 +16,12 @@ escape_like
 .. autofunction:: escape_like
 
 
+get_columns
+^^^^^^^^^^^
+
+.. autofunction:: get_columns
+
+
 query_entities
 ^^^^^^^^^^^^^^
 
@@ -40,6 +46,12 @@ naturally_equivalent
 .. autofunction:: naturally_equivalent
 
 
+primary_keys
+^^^^^^^^^^^^
+
+.. autofunction:: primary_keys
+
+
 sort_query
 ^^^^^^^^^^
 

sqlalchemy_utils/functions/orm.py

@@ -14,21 +14,42 @@ from sqlalchemy.orm.query import _ColumnEntity
 from sqlalchemy.orm.util import AliasedInsp
 
 
-def primary_keys(obj_or_class):
+def primary_keys(mixed):
     """
-    Return an OrderedDict of all primary keys for given declarative class or
-    object.
-    """
-    if not isclass(obj_or_class):
-        obj_or_class = obj_or_class.__class__
+    Return an OrderedDict of all primary keys for given Table object,
+    declarative class or declarative class instance.
 
+    :param mixed:
+        SA Table object, SA declarative class or SA declarative class instance
+
+    .. versionchanged: 0.25.3
+        Made the function return an ordered dictionary instead of generator.
+        This change was made to support primary key aliases.
+
+    .. seealso:: :func:`get_columns`
+    """
     columns = OrderedDict()
-    for key, column in sa.inspect(obj_or_class).columns.items():
+    for key, column in get_columns(mixed).items():
         if column.primary_key:
             columns[key] = column
     return columns
 
 
+def get_columns(mixed):
+    """
+    Return a collection of all Column objects for given SQLAlchemy
+    Table object, declarative class or declarative class instance.
+
+    :param mixed:
+        SA Table object, SA declarative class or SA declarative class instance
+    """
+    if isinstance(mixed, sa.Table):
+        return mixed.c
+    if not isclass(mixed):
+        mixed = mixed.__class__
+    return sa.inspect(mixed).columns
+
+
 def table_name(obj):
     """
     Return table name of given target, declarative class or the

tests/functions/test_get_columns.py

@@ -0,0 +1,31 @@
+import sqlalchemy as sa
+from sqlalchemy_utils.functions.orm import get_columns
+from tests import TestCase
+
+
+class TestGetColumns(TestCase):
+    def create_models(self):
+        class Building(self.Base):
+            __tablename__ = 'building'
+            id = sa.Column('_id', sa.Integer, primary_key=True)
+            name = sa.Column('_name', sa.Unicode(255))
+
+        self.Building = Building
+
+    def test_table(self):
+        assert isinstance(
+            get_columns(self.Building.__table__),
+            sa.sql.base.ImmutableColumnCollection
+        )
+
+    def test_declarative_class(self):
+        assert isinstance(
+            get_columns(self.Building),
+            sa.util._collections.OrderedProperties
+        )
+
+    def test_declarative_object(self):
+        assert isinstance(
+            get_columns(self.Building()),
+            sa.util._collections.OrderedProperties
+        )