From cd92d2b1ff40db30db35f21f06a31c51f9b900a3 Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Tue, 24 Sep 2013 17:53:18 +0300
Subject: [PATCH] Made defer_except to allow class attributes

---
 sqlalchemy_utils/functions/__init__.py     | 28 ++--------------
 sqlalchemy_utils/functions/defer_except.py | 39 ++++++++++++++++++++++
 tests/test_defer_except.py                 | 14 ++++++++
 tests/test_utility_functions.py            |  7 ----
 4 files changed, 55 insertions(+), 33 deletions(-)
 create mode 100644 sqlalchemy_utils/functions/defer_except.py
 create mode 100644 tests/test_defer_except.py

diff --git a/sqlalchemy_utils/functions/__init__.py b/sqlalchemy_utils/functions/__init__.py
index 9b191fc..cb14f81 100644
--- a/sqlalchemy_utils/functions/__init__.py
+++ b/sqlalchemy_utils/functions/__init__.py
@@ -2,16 +2,16 @@ from collections import defaultdict
 import six
 import datetime
 import sqlalchemy as sa
-from sqlalchemy.orm import defer
-from sqlalchemy.orm.properties import ColumnProperty
 from sqlalchemy.orm.query import Query
 from sqlalchemy.schema import MetaData, Table, ForeignKeyConstraint
 from .batch_fetch import batch_fetch, with_backrefs, CompositePath
+from .defer_except import defer_except
 from .sort_query import sort_query, QuerySorterException
 
 
 __all__ = (
     batch_fetch,
+    defer_except,
     sort_query,
     with_backrefs,
     CompositePath,
@@ -19,30 +19,6 @@ __all__ = (
 )
 
 
-def defer_except(query, columns):
-    """
-    Deferred loads all columns in given query, except the ones given.
-
-    This function is very useful when working with models with myriad of
-    columns and you want to deferred load many columns.
-
-        >>> from sqlalchemy_utils import defer_except
-        >>> query = session.query(Article)
-        >>> query = defer_except(Article, [Article.id, Article.name])
-
-    :param columns: columns not to deferred load
-    """
-    model = query._entities[0].entity_zero.class_
-    fields = set(model._sa_class_manager.values())
-    for field in fields:
-        property_ = field.property
-        if isinstance(property_, ColumnProperty):
-            column = property_.columns[0]
-            if column.name not in columns:
-                query = query.options(defer(property_.key))
-    return query
-
-
 def escape_like(string, escape_char='*'):
     """
     Escapes the string paremeter used in SQL LIKE expressions
diff --git a/sqlalchemy_utils/functions/defer_except.py b/sqlalchemy_utils/functions/defer_except.py
new file mode 100644
index 0000000..3d89dd2
--- /dev/null
+++ b/sqlalchemy_utils/functions/defer_except.py
@@ -0,0 +1,39 @@
+import six
+from sqlalchemy import inspect
+from sqlalchemy.orm import defer
+from sqlalchemy.orm.properties import ColumnProperty
+
+
+def property_names(properties):
+    names = []
+    for property_ in properties:
+        if isinstance(property_, six.string_types):
+            names.append(property_)
+        else:
+            names.append(property_.key)
+    return names
+
+
+def defer_except(query, properties):
+    """
+    Deferred loads all properties in given query, except the ones given.
+
+    This function is very useful when working with models with myriad of
+    properties and you want to deferred load many properties.
+
+        >>> from sqlalchemy_utils import defer_except
+        >>> query = session.query(Article)
+        >>> query = defer_except(Article, [Article.id, Article.name])
+
+    :param query: SQLAlchemy Query object to apply the deferred loading to
+    :param properties: properties not to deferred load
+    """
+
+    allowed_names = property_names(properties)
+
+    model = query._entities[0].entity_zero.class_
+    for property_ in inspect(model).attrs:
+        if isinstance(property_, ColumnProperty):
+            if property_.key not in allowed_names:
+                query = query.options(defer(property_.key))
+    return query
diff --git a/tests/test_defer_except.py b/tests/test_defer_except.py
new file mode 100644
index 0000000..f04543b
--- /dev/null
+++ b/tests/test_defer_except.py
@@ -0,0 +1,14 @@
+from sqlalchemy_utils import defer_except
+from tests import TestCase
+
+
+class TestDeferExcept(TestCase):
+    def test_supports_properties_as_strings(self):
+        query = self.session.query(self.Article)
+        query = defer_except(query, ['id'])
+        assert str(query) == 'SELECT article.id AS article_id \nFROM article'
+
+    def test_supports_properties_as_class_attributes(self):
+        query = self.session.query(self.Article)
+        query = defer_except(query, [self.Article.id])
+        assert str(query) == 'SELECT article.id AS article_id \nFROM article'
diff --git a/tests/test_utility_functions.py b/tests/test_utility_functions.py
index 1f1c130..f6d7bab 100644
--- a/tests/test_utility_functions.py
+++ b/tests/test_utility_functions.py
@@ -12,13 +12,6 @@ class TestEscapeLike(TestCase):
         assert escape_like('_*%') == '*_***%'
 
 
-class TestDeferExcept(TestCase):
-    def test_deferred_loads_all_columns_except_the_ones_given(self):
-        query = self.session.query(self.Article)
-        query = defer_except(query, ['id'])
-        assert str(query) == 'SELECT article.id AS article_id \nFROM article'
-
-
 class TestFindNonIndexedForeignKeys(TestCase):
     # dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'