Change Data class and add DataElement

Author: seignovert

opus/data.py

@@ -4,30 +4,32 @@ class Data(object):
     def __init__(self, json):
         self.json = json
         self.index = 0
+        self._data = {}
+        for row in self.json['page']:
+            el = DataElement(self.columns, row)
+            self._data[str(el)] = el
 
     def __repr__(self):
-        return 'OPUS API Data object (with {} elements)'.format(self.count)
+        return 'OPUS API Data object (with {} elements):\n'.format(len(self)) + \
+               '\n'.join(' - {}'.format(key) for key, _ in self.items())
 
     def __len__(self):
-        return self.count
+        return len(self._data)
 
-    def __getitem__(self, index):
-        return self.json['page'][index]
+    def __getitem__(self, attr):
+        return self._data[attr]
 
     def __iter__(self):
-        return self
+        return iter(self._data)
 
-    def __next__(self):
-        try:
-            result = self.__getitem__(self.index)
-        except IndexError:
-            self.index = 0
-            raise StopIteration
-        self.index += 1
-        return result
+    def keys(self):
+        return self._data.keys()
 
-    def next(self):
-        return self.__next__()
+    def items(self):
+        return self._data.items()
+
+    def values(self):
+        return self._data.values()
 
     @property
     def count(self):
@@ -48,3 +50,47 @@ def page_no(self):
     @property
     def labels(self):
         return self.json['labels']
+
+    @property
+    def columns(self):
+        return self.json['columns']
+
+class DataElement(object):
+    def __init__(self, columns, row):
+        self._columns = columns
+        self._row = row
+        self._id = None
+        self._data = {}
+        for column, value in zip(columns, row):
+            if 'Ring Observation ID' in column:
+                self._id = value
+            else:
+                self._data[column] = value
+
+        if self._id is None:
+            raise KeyError('`Ring Observation ID` key is missing in `cols` query')
+
+    def __repr__(self):
+        return 'Ring Observation ID: {}\n'.format(str(self)) + \
+               '\n'.join(' - {}: {}'.format(key, value) for key, value in self.items())
+
+    def __str__(self):
+        return self._id
+
+    def __len__(self):
+        return len(self._data)
+
+    def __getitem__(self, attr):
+        return self._data[attr]
+
+    def __iter__(self):
+        return iter(self._data)
+
+    def keys(self):
+        return self._data.keys()
+
+    def items(self):
+        return self._data.items()
+
+    def values(self):
+        return self._data.values()

tests/test_data.py

@@ -3,7 +3,7 @@
 
 import json as JSON
 
-from opus.data import Data
+from opus.data import Data, DataElement
 
 @pytest.fixture
 def json():
@@ -13,9 +13,16 @@ def json():
 def data(json):
     return Data(json)
 
+@pytest.fixture
+def el(data):
+    return data['S_IMG_CO_ISS_1508094647_N']
+
 def test_data_meta(data):
-    assert repr(data) == 'OPUS API Data object (with 10 elements)'
+    r = repr(data)
+    assert 'OPUS API Data object (with 10 elements)' in r
+    assert '- S_IMG_CO_ISS_1508094647_N' in r
     assert len(data) == 10
+    assert data.count == 10
     assert data.limit == 10
     assert data.page_no == 2
     assert data.order == 'obs_general.time1'
@@ -32,18 +39,33 @@ def test_data_labels(data):
     assert labels[6] == 'Observation Stop Time (UTC)'
 
 
-def test_data_images(data):
-    img = data[0]
-    assert img[0] == 'S_IMG_CO_ISS_1508094647_N'
-    assert img[1] == 'SAT'
-    assert img[2] == 'PAN'
-    assert img[3] == 87.786
-    assert img[4] == 88.169
-    assert img[5] == '2005-288T18:42:08.622'
-    assert img[6] == '2005-288T18:42:09.302'
+def test_data_iter(data):
+    for key in data:
+        assert key in data.keys()
+        break
 
+    for value in data.values():
+        assert isinstance(value, DataElement)
+        break
 
-def test_data_iter(data):
-    data.next() == data[0]
-    for ii, img in enumerate(data):
-        img = data[ii]
+
+def test_data_element(el):
+    assert el['Planet'] == 'SAT'
+    assert el['Intended Target Name'] == 'PAN'
+    assert el['Observed Phase Angle (Min)'] == 87.786
+    assert el['Observed Phase Angle (Max)'] == 88.169
+    assert el['Observation Start Time (UTC)'] == '2005-288T18:42:08.622'
+    assert el['Observation Stop Time (UTC)'] == '2005-288T18:42:09.302'
+    assert len(el) == 6
+    assert 'Planet' in repr(el)
+    assert 'SAT' in el.values()
+
+
+def test_data_element_iter(el):
+    for key in el:
+        assert key in el.keys()
+        break
+
+def test_data_element_err():
+    with pytest.raises(KeyError):
+        DataElement(['abc'], ['foo'])