From dd321e3a9ce1baf97351d14cde681e4e419515ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20P=C3=A9ters?= <fpeters@entrouvert.com>
Date: Fri, 12 Apr 2019 15:21:29 +0200
Subject: [PATCH 2/3] lazy: add queryset-like features to lazy formdef (#32244)

{% for formdata in form.formdef.objects.pending %}
  {{formdata.number}}
{% endfor %}
---
 tests/test_formdata.py | 45 +++++++++++++++++++++++++++++
 wcs/qommon/storage.py  |  5 ++++
 wcs/variables.py       | 64 ++++++++++++++++++++++++++++++++++++++----
 3 files changed, 109 insertions(+), 5 deletions(-)

diff --git a/tests/test_formdata.py b/tests/test_formdata.py
index eb13f2c5f..8b13e12ca 100644
--- a/tests/test_formdata.py
+++ b/tests/test_formdata.py
@@ -16,6 +16,7 @@ from wcs.conditions import Condition
 from wcs.formdef import FormDef
 from wcs.formdata import Evolution
 from wcs.roles import Role
+from wcs import sessions
 from wcs.variables import LazyFormData
 from wcs.workflows import Workflow, WorkflowStatusItem, WorkflowCriticalityLevel, WorkflowBackofficeFieldsFormDef
 from wcs.wf.anonymise import AnonymiseWorkflowStatusItem
@@ -616,6 +617,7 @@ def variable_test_data(pub):
     pub.substitutions.feed(pub)
     pub.substitutions.feed(formdef)
     pub.substitutions.feed(formdata)
+    return LazyFormData(formdata)
 
 def test_lazy_formdata(pub, variable_test_data):
     formdata = FormDef.select()[0].data_class().select()[0]
@@ -740,6 +742,49 @@ def test_lazy_formdata(pub, variable_test_data):
     lazy_formdata = LazyFormData(formdata)
     assert lazy_formdata.tracking_code == tracking_code.id
 
+def test_lazy_formdata_queryset(pub, variable_test_data):
+    lazy_formdata = variable_test_data
+    data_class = lazy_formdata._formdef.data_class()
+    for i in range(6):
+        formdata = data_class()
+        formdata.just_created()
+        formdata.store()
+    for i in range(4):
+        formdata = data_class()
+        formdata.just_created()
+        formdata.jump_status('finished')
+        formdata.store()
+
+    formdata = data_class()
+    formdata.status = 'draft'
+    formdata.store()
+
+    assert lazy_formdata.objects.count() == 11
+    assert lazy_formdata.objects.drafts().count() == 1
+    assert lazy_formdata.objects.pending().count() == 7
+    assert lazy_formdata.objects.done().count() == 4
+    assert lazy_formdata.objects.drafts().count() == 1
+    # check __len__
+    assert len(lazy_formdata.objects.drafts()) == 1
+
+    assert lazy_formdata.objects.current_user().count() == 0
+    pub.get_request()._user = ()  # reset cache
+    pub.get_request().session = sessions.BasicSession(id=1)
+    pub.get_request().session.set_user(pub.user_class.select()[0].id)
+    assert lazy_formdata.objects.current_user().count() == 1
+
+    qs = lazy_formdata.objects.drafts()
+    # check __iter__
+    for draft in qs:
+        assert draft.internal_id == formdata.id
+    # check __getitem__
+    assert qs[0].internal_id == formdata.id
+    # check against for cached resultset
+    assert qs[0].internal_id == formdata.id
+    # check __iter__ with cached resultset
+    for draft in qs:
+        assert draft.internal_id == formdata.id
+
 
 def test_lazy_variables(pub, variable_test_data):
     formdata = FormDef.select()[0].data_class().select()[0]
diff --git a/wcs/qommon/storage.py b/wcs/qommon/storage.py
index 8bfac1136..4e0d59a2b 100644
--- a/wcs/qommon/storage.py
+++ b/wcs/qommon/storage.py
@@ -289,6 +289,11 @@ class StorableObject(object):
             objects = _take(objects, limit, offset)
         return list(objects)
 
+    @classmethod
+    def select_iterator(cls, **kwargs):
+        for obj in cls.select(**kwargs):
+            yield obj
+
     @classmethod
     def has_key(cls, id):
         filename = os.path.join(cls.get_objects_dir(), fix_key(id))
diff --git a/wcs/variables.py b/wcs/variables.py
index 732f2c1da..31ae3c775 100644
--- a/wcs/variables.py
+++ b/wcs/variables.py
@@ -16,27 +16,81 @@
 
 from django.utils.functional import SimpleLazyObject
 
-from quixote import get_publisher
+from quixote import get_publisher, get_request
 
 import qommon.misc
 from qommon.evalutils import make_datetime
 from qommon.templatetags.qommon import parse_datetime
+from qommon.storage import (Or, Equal, NotEqual)
 
 
 class LazyFormDefObjectsManager(object):
-    def __init__(self, formdef):
+    def __init__(self, formdef, criterias=None):
         self._formdef = formdef
+        if criterias is None:
+            criterias = [NotEqual('status', 'draft')]
+            # add custom marker to criteria so it can be found back and removed in
+            # drafts().
+            criterias[0].exclude_drafts = True
+        self._criterias = criterias
+        self._cached_resultset = None
 
-    @property
     def count(self):
-        return self._formdef.data_class().count()
+        return self._formdef.data_class().count(clause=self._criterias)
+
+    def all(self):
+        # (mimicking django)
+        return LazyFormDefObjectsManager(self._formdef)
+
+    def pending(self):
+        status_filters = ['wf-%s' % x.id for x in
+                          self._formdef.workflow.get_not_endpoint_status()]
+        criterias = [Or([Equal('status', x) for x in status_filters])]
+        return LazyFormDefObjectsManager(
+                self._formdef, self._criterias + criterias)
+
+    def current_user(self):  # filter on current user
+        user = get_request().user
+        return LazyFormDefObjectsManager(self._formdef,
+                self._criterias + [Equal('user_id', str(user.id) if user else '-1')])
+
+    def drafts(self):
+        criterias = [x for x in self._criterias if not getattr(x, 'exclude_drafts', False)]
+        return LazyFormDefObjectsManager(self._formdef,
+                criterias + [Equal('status', 'draft')])
+
+    def done(self):
+        status_filters = ['wf-%s' % x.id for x in
+                          self._formdef.workflow.get_endpoint_status()]
+        criterias = [Or([Equal('status', x) for x in status_filters])]
+        return LazyFormDefObjectsManager(self._formdef,
+                self._criterias + criterias)
 
     def __getattr__(self, attribute):
+        # backward compatibility
         if attribute.startswith('count_status_'):
             status = attribute[len('count_status_'):]
             return len(self._formdef.data_class().get_ids_with_indexed_value(
                 'status', 'wf-%s' % status))
-        return super(FormDefSubstVar, self).__getattr__(attribute)
+        return super(LazyFormDefObjectsManager, self).__getattr__(attribute)
+
+    def __len__(self):
+        if self._cached_resultset:
+            return len(self._cached_resultset)
+        return self.count()
+
+    def __getitem__(self, key):
+        if self._cached_resultset is None:
+            self._cached_resultset = self._formdef.data_class().select(clause=self._criterias)
+        return LazyFormData(self._cached_resultset[key])
+
+    def __iter__(self):
+        if self._cached_resultset:
+            for formdata in self._cached_resultset:
+                yield LazyFormData(formdata)
+        data_class = self._formdef.data_class()
+        for formdata in self._formdef.data_class().select_iterator(clause=self._criterias):
+            yield LazyFormData(formdata)
 
 
 class LazyFormDef(object):
-- 
2.20.1

