Improved queryset getter for use .count instead of len

Author: syrusakbary

graphene/contrib/django/utils.py

@@ -1,5 +1,8 @@
 from django.db import models
 from django.db.models.manager import Manager
+from django.db.models.query import QuerySet
+
+from graphene.utils import LazyList
 
 
 def get_type_for_model(schema, model):
@@ -19,7 +22,18 @@ def get_reverse_fields(model):
             yield related
 
 
+class WrappedQueryset(LazyList):
+
+    def __len__(self):
+        # Dont calculate the length using len(queryset), as this will
+        # evaluate the whole queryset and return it's length.
+        # Use .count() instead
+        return self._origin.count()
+
+
 def maybe_queryset(value):
     if isinstance(value, Manager):
         value = value.get_queryset()
+    if isinstance(value, QuerySet):
+        return WrappedQueryset(value)
     return value

graphene/relay/types.py

@@ -4,6 +4,7 @@
 from functools import wraps
 
 import six
+
 from graphql_relay.connection.arrayconnection import connection_from_list
 from graphql_relay.node.node import to_global_id
 

graphene/utils/__init__.py

@@ -3,8 +3,9 @@
 from .caching import cached_property, memoize
 from .misc import enum_to_graphql_enum
 from .resolve_only_args import resolve_only_args
+from .lazylist import LazyList
 
 
 __all__ = ['to_camel_case', 'to_snake_case', 'ProxySnakeDict',
            'cached_property', 'memoize', 'enum_to_graphql_enum',
-           'resolve_only_args']
+           'resolve_only_args', 'LazyList']

graphene/utils/lazylist.py

@@ -0,0 +1,43 @@
+class LazyList(object):
+
+    def __init__(self, origin, state=None):
+        self._origin = origin
+        self._state = state or []
+        self._origin_iter = None
+        self._finished = False
+
+    def __iter__(self):
+        return self if not self._finished else iter(self._state)
+
+    def iter(self):
+        return self.__iter__()
+
+    def __len__(self):
+        return self._origin.__len__()
+
+    def __next__(self):
+        try:
+            if not self._origin_iter:
+                self._origin_iter = self._origin.__iter__()
+            n = next(self._origin_iter)
+        except StopIteration as e:
+            self._finished = True
+            raise e
+        else:
+            self._state.append(n)
+            return n
+
+    def next(self):
+        return self.__next__()
+
+    def __getitem__(self, key):
+        item = self._origin[key]
+        if isinstance(key, slice):
+            return self.__class__(item)
+        return item
+
+    def __getattr__(self, name):
+        return getattr(self._origin, name)
+
+    def __repr__(self):
+        return "<{} {}>".format(self.__class__.__name__, repr(self._origin))

graphene/utils/tests/test_lazylist.py

@@ -0,0 +1,23 @@
+from py.test import raises
+
+from ..lazylist import LazyList
+
+
+def test_lazymap():
+    data = list(range(10))
+    lm = LazyList(data)
+    assert len(lm) == 10
+    assert lm[1] == 1
+    assert isinstance(lm[1:4], LazyList)
+    assert lm.append == data.append
+    assert repr(lm) == '<LazyList [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]>'
+
+
+def test_lazymap_iter():
+    data = list(range(2))
+    lm = LazyList(data)
+    iter_lm = iter(lm)
+    assert iter_lm.next() == 0
+    assert iter_lm.next() == 1
+    with raises(StopIteration):
+        iter_lm.next()