From 18a40f68f35f8b0c958a15fd1433033d336de149 Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Mon, 21 Apr 2014 10:57:14 +0300
Subject: [PATCH] Add get_columns utility function

---
 CHANGES.rst                         |  7 ++++++
 docs/model_helpers.rst              | 12 ++++++++++
 sqlalchemy_utils/functions/orm.py   | 35 +++++++++++++++++++++++------
 tests/functions/test_get_columns.py | 31 +++++++++++++++++++++++++
 4 files changed, 78 insertions(+), 7 deletions(-)
 create mode 100644 tests/functions/test_get_columns.py

diff --git a/CHANGES.rst b/CHANGES.rst
index f443971..032eb46 100644
--- a/CHANGES.rst
+++ b/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)
 ^^^^^^^^^^^^^^^^^^^
 
diff --git a/docs/model_helpers.rst b/docs/model_helpers.rst
index 9fbf30d..43ae050 100644
--- a/docs/model_helpers.rst
+++ b/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
 ^^^^^^^^^^
 
diff --git a/sqlalchemy_utils/functions/orm.py b/sqlalchemy_utils/functions/orm.py
index 1a276d7..12fd004 100644
--- a/sqlalchemy_utils/functions/orm.py
+++ b/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
diff --git a/tests/functions/test_get_columns.py b/tests/functions/test_get_columns.py
new file mode 100644
index 0000000..270486a
--- /dev/null
+++ b/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
+        )