Merge pull request #774 from gucky92/fetch_format_bug

fix when fetch_format='frame'

Author: guzman-raphael

datajoint/fetch.py

@@ -173,7 +173,9 @@ def __call__(self, *attrs, offset=None, limit=None, order_by=None, format=None,
             attributes = [a for a in attrs if not is_key(a)]
             ret = self._expression.proj(*attributes).fetch(
                 offset=offset, limit=limit, order_by=order_by,
-                as_dict=False, squeeze=squeeze, download_path=download_path)
+                as_dict=False, squeeze=squeeze, download_path=download_path,
+                format='array'
+            )
             if attrs_as_dict:
                 ret = [{k: v for k, v in zip(ret.dtype.names, x) if k in attrs} for x in ret]
             else:

tests/test_fetch.py

@@ -220,3 +220,27 @@ def test_nullable_numbers(self):
         assert_true(None in i)
         assert_true(any(np.isnan(d)))
         assert_true(any(np.isnan(f)))
+
+    def test_fetch_format(self):
+        """test fetch_format='frame'"""
+        dj.config['fetch_format'] = 'frame'
+        # test if lists are both dicts
+        list1 = sorted(self.subject.proj().fetch(as_dict=True), key=itemgetter('subject_id'))
+        list2 = sorted(self.subject.fetch(dj.key), key=itemgetter('subject_id'))
+        for l1, l2 in zip(list1, list2):
+            assert_dict_equal(l1, l2, 'Primary key is not returned correctly')
+
+        # tests if pandas dataframe
+        tmp = self.subject.fetch(order_by='subject_id')
+        assert_true(isinstance(tmp, pandas.DataFrame))
+        tmp = tmp.to_records()
+
+        subject_notes, key, real_id = self.subject.fetch('subject_notes', dj.key, 'real_id')
+
+        np.testing.assert_array_equal(sorted(subject_notes), sorted(tmp['subject_notes']))
+        np.testing.assert_array_equal(sorted(real_id), sorted(tmp['real_id']))
+        list1 = sorted(key, key=itemgetter('subject_id'))
+        for l1, l2 in zip(list1, list2):
+            assert_dict_equal(l1, l2, 'Primary key is not returned correctly')
+        # revert configuration of fetch format
+        dj.config['fetch_format'] = 'array'