Merge pull request #388 from Crunch-io/set-date-field-on-categories

Add support for editing the `date` field on categoricals
Crunch-io · Oct 13, 2020 · 26fd48e · 26fd48e
2 parents 7aafbe4 + 5bab87e
commit 26fd48e
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 6 deletions.
diff --git a/scrunch/categories.py b/scrunch/categories.py
@@ -6,7 +6,7 @@
 
 
 class Category(ReadOnly):
-    _MUTABLE_ATTRIBUTES = {'name', 'numeric_value', 'missing', 'selected'}
+    _MUTABLE_ATTRIBUTES = {'name', 'numeric_value', 'missing', 'selected', 'date'}
     _IMMUTABLE_ATTRIBUTES = {'id'}
     _ENTITY_ATTRIBUTES = _MUTABLE_ATTRIBUTES | _IMMUTABLE_ATTRIBUTES
 
@@ -21,16 +21,23 @@ def __getattr__(self, item):
                 return self._category.get('selected', False)
             return self._category[item]  # Has to exist
 
-        # Attribute doesn't exists, must raise an AttributeError
-        raise AttributeError('Category %s has no attribute %s' % (
-            self.category['name'], item))
+        # Attribute doesn't exist, must raise an AttributeError
+        raise AttributeError(
+            'Category %s has no attribute %s' % (self.category['name'], item)
+        )
 
     def __repr__(self):
         attrs = self.as_dict()
         return 'Category(%s)' % ', '.join('%s=%s' % c for c in attrs.items())
 
     def as_dict(self, **kwargs):
-        dct = {attr: getattr(self, attr) for attr in self._ENTITY_ATTRIBUTES}
+        attributes = self._ENTITY_ATTRIBUTES - {'date'}  # `date` needs special handling
+        dct = {attr: getattr(self, attr) for attr in attributes}
+        try:
+            dct['date'] = getattr(self, 'date')
+        except KeyError:
+            # `date` is not there, just move on
+            pass
         if PY2:
             dct['name'] = dct['name'].encode("ascii", "replace")
         dct.update(**kwargs or {})

diff --git a/scrunch/tests/integration/scrunch_workflow_integration_test.py b/scrunch/tests/integration/scrunch_workflow_integration_test.py
@@ -9,6 +9,12 @@
 
 pytest.mark.skip('skip test discovery on this module')
 
+try:
+    from pycrunch import pandaslib
+except ImportError:
+    # pandas is not installed
+    pandaslib = None
+
 from scrunch import connect
 from scrunch.datasets import Variable, get_geodata
 from scrunch.streaming_dataset import StreamingDataset
@@ -412,6 +418,7 @@ def isnan(obj):
 
 def main():
     assert not invalid_credentials()
+    assert pandaslib, 'Pandas library not installed'
 
     # Login.
     site = connect(CRUNCH_USER, CRUNCH_PASSWORD, CRUNCH_URL)

diff --git a/scrunch/tests/test_categories.py b/scrunch/tests/test_categories.py
@@ -12,6 +12,13 @@ def TEST_CATEGORIES():
         {"id": -1, "name": "No Data", "missing": True, "numeric_value": None}
     ]
 
+def TEST_CATEGORIES_WITH_DATE():
+    return [
+        {"id": 1, "name": "Female", "missing": False, "numeric_value": None, "date": "2020-01-01"},
+        {"id": 2, "name": "Male", "missing": False, "numeric_value": None, "date": "2020-02-02"},
+        {"id": -1, "name": "No Data", "missing": True, "numeric_value": None}
+    ]
+
 
 class EditableMock(MagicMock):
     def edit(self, **kwargs):
@@ -82,7 +89,7 @@ def test_delete_category(self):
             {'numeric_value': None, 'missing': True, 'id': -1, 'name': 'No Data'}
         ])
 
-    def test_Category_attribute_writes(self):
+    def test_category_attribute_writes(self):
         resource = EditableMock()
         resource.entity.body = dict(
             categories=TEST_CATEGORIES(),
@@ -151,6 +158,35 @@ def test_edit_derived(self):
             {'numeric_value': None, 'missing': True, 'id': -1, 'name': 'No Data'}
         ])
 
+    def test_read_category_date(self):
+        resource = EditableMock()
+        resource.entity.body = dict(
+            categories=TEST_CATEGORIES_WITH_DATE(),
+            type='categorical'
+        )
+        variable = Variable(resource, MagicMock())
+        self.assertEqual(variable.categories[1].date, '2020-01-01')
+        self.assertEqual(variable.categories[2].date, '2020-02-02')
+        with self.assertRaises(KeyError):
+            # The `No Data` category doesn't provide a `date field
+            _ = variable.categories[3].date
+
+    def test_edit_category_date(self):
+        resource = EditableMock()
+        resource.entity.body = dict(
+            categories=TEST_CATEGORIES_WITH_DATE(),
+            type='categorical'
+        )
+        variable = Variable(resource, MagicMock())
+        variable.categories[1].edit(date='2021-01-01')
+        resource.entity._edit.assert_called_with(categories=[
+            {'numeric_value': None, 'selected': False, 'id': 1, 'missing': False, 'name': 'Female', 'date': '2021-01-01'},
+            {'numeric_value': None, 'missing': False, 'id': 2, 'name': 'Male', 'date': '2020-02-02'},
+            {'numeric_value': None, 'missing': True, 'id': -1, 'name': 'No Data'}
+        ])
+        resource.entity.refresh.assert_called_once()
+        self.assertEqual(variable.categories[1].date, '2021-01-01')
+
 
 class TestCategoryList(TestCase):
     def test_reorder(self):