Add class argument support for identity

CHANGES.rst

@@ -4,6 +4,14 @@ Changelog
 Here you can see the full list of changes between each SQLAlchemy-Utils release.
 
 
+
+0.24.1 (2014-02-21)
+^^^^^^^^^^^^^^^^^^^
+
+- Made identity return a tuple in all cases
+- Added support for declarative model classes as identity function's first argument
+
+
 0.24.0 (2014-02-18)
 ^^^^^^^^^^^^^^^^^^^
 

setup.py

@@ -42,7 +42,7 @@ for name, requirements in extras_require.items():
 
 setup(
     name='SQLAlchemy-Utils',
-    version='0.24.0',
+    version='0.24.1',
     url='https://github.com/kvesteri/sqlalchemy-utils',
     license='BSD',
     author='Konsta Vesterinen, Ryan Leckey, Janne Vanhala, Vesa Uimonen',

sqlalchemy_utils/__init__.py

@@ -56,7 +56,7 @@ from .types import (
 )
 
 
-__version__ = '0.24.0'
+__version__ = '0.24.1'
 
 
 __all__ = (

sqlalchemy_utils/functions/orm.py

@@ -1,4 +1,5 @@
 from functools import partial
+from operator import attrgetter
 from toolz import curry, first
 import six
 import sqlalchemy as sa
@@ -220,9 +221,6 @@ def declarative_base(model):
     return model
 
 
-from operator import attrgetter
-
-
 def getdotattr(obj_or_class, dot_path):
     """
     Allow dot-notated strings to be passed to `getattr`.
@@ -293,12 +291,13 @@ def has_changes(obj, attr):
     )
 
 
-def identity(obj):
+def identity(obj_or_class):
     """
-    Return the identity of given sqlalchemy declarative model instance as a
-    tuple. This differs from obj._sa_instance_state.identity in a way that it
-    always returns the identity even if object is still in transient state (
-    new object that is not yet persisted into database).
+    Return the identity of given sqlalchemy declarative model class or instance
+    as a tuple. This differs from obj._sa_instance_state.identity in a way that
+    it always returns the identity even if object is still in transient state (
+    new object that is not yet persisted into database). Also for classes it
+    returns the identity attributes.
 
     ::
 
@@ -322,19 +321,19 @@ def identity(obj):
         inspect(user).identity  # (1, )
 
 
+    You can also use identity for classes::
+
+
+        identity(User)  # (User.id, )
+
     .. versionadded: 0.21.0
 
     :param obj: SQLAlchemy declarative model object
     """
-    id_ = []
-    for column in sa.inspect(obj.__class__).columns:
-        if column.primary_key:
-            id_.append(getattr(obj, column.name))
-
-    if all(value is None for value in id_):
-        return None
-    else:
-        return tuple(id_)
+    return tuple(
+        getattr(obj_or_class, column.name)
+        for column in primary_keys(obj_or_class)
+    )
 
 
 def naturally_equivalent(obj, obj2):

tests/functions/test_identity.py

@@ -13,7 +13,7 @@ class TestIdentity(TestCase):
         self.Building = Building
 
     def test_for_transient_class_without_id(self):
-        assert identity(self.Building()) is None
+        assert identity(self.Building()) == (None, )
 
     def test_for_transient_class_with_id(self):
         building = self.Building(name=u'Some building')
@@ -21,3 +21,6 @@ class TestIdentity(TestCase):
         self.session.flush()
 
         assert identity(building) == (building.id, )
+
+    def test_identity_for_class(self):
+        assert identity(self.Building) == (self.Building.id, )