Add del helper function and __delitem__ to Data (#208)

* Add del helper function and __delitem__ to Data

* Tests for deleting fields/trials

* Fixed incorrect var name

* Update test_data_object.py

Fixing test variable names

* Fix expected result of test

* Update data.rst

* Update naplib/data.py

* Update naplib/data.py

* Update __init__.py

---------

Co-authored-by: Gavin Mischler <[email protected]>

docs/references/data.rst

@@ -8,8 +8,8 @@ Data object
 
 .. autoclass:: Data
     :members:
-    :special-members: __getitem__, __setitem__, __len__
-    :exclude-members: set_field, get_field
+    :special-members: __getitem__, __setitem__, __len__, __delitem__
+    :exclude-members: set_field, get_field, delete_field
 
 .. minigallery:: naplib.Data
         :add-heading: Examples using ``Data ``

naplib/__init__.py

@@ -56,5 +56,5 @@ def set_logging(level: Union[int, str]):
 from .data import Data, join_fields, concat
 import naplib.naplab
 
-__version__ = "2.1.0"
+__version__ = "2.2.0"
 

naplib/data.py

@@ -136,6 +136,20 @@ def set_field(self, fielddata, fieldname):
             raise Exception(f'Length of field ({len(fielddata)}) is not equal to length of this Data ({len(self)})')
         for i, trial in enumerate(self.data):
             trial[fieldname] = fielddata[i]
+
+    def delete_field(self, fieldname):
+        '''
+        Remove an entire field from the Data object.
+
+        Parameters
+        ----------
+        fieldname : string
+            Name of field to delete.
+        '''
+        if not isinstance(fieldname, str):
+            raise TypeError(f'Field must be a str, but found {type(fieldname)}')
+        for trial in self.data:
+            del trial[fieldname]
 
     def get_field(self, fieldname):
         '''
@@ -281,6 +295,48 @@ def __setitem__(self, index, data):
             else:
                 self.data[index] = data
 
+    def __delitem__(self, index):
+        '''
+        Delete a specific trial or set of trials, or delete a specific field, using
+        bracket indexing. See examples below for details.
+
+        Parameters
+        ----------
+        index : int or string
+            Which trial to delete, or which field to delete. If an integer, must be 
+            < the length of the Data, since you can only delete an existing trial 
+
+        Examples
+        --------
+        >>> # Delete a field of a Data
+        >>> from naplib import Data
+        >>> trial_data = [{'name': 'Zero', 'trial': 0, 'resp': [[0,1],[2,3]]},
+        ...               {'name': 'One', 'trial': 1, 'resp': [[4,5],[6,7]]}]
+        >>> data = Data(trial_data)
+        >>> del data[0]
+        >>> data[0]
+        {'name': 'One', 'trial': 1, 'resp': [[4, 5], [6, 7]]}
+
+        >>> # We can also delete all values of a field across trials
+        >>> trial_data = [{'name': 'Zero', 'trial': 0, 'resp': [[0,1],[2,3]]},
+        ...               {'name': 'One', 'trial': 1, 'resp': [[4,5],[6,7]]}]
+        >>> data = Data(trial_data)
+        >>> del data['name']
+        >>> data[0]
+        {'trial': 0, 'resp': [[0, 1], [2, 3]]}
+        '''
+        if isinstance(index, str):
+            self.delete_field(index)
+        elif isinstance(index, int):
+            if index >= len(self):
+                raise IndexError((f'Index is too large. Current data is length {len(self)} '
+                    'but tried to delete index {index}. If you want to add to the end of the list '
+                    'of trials, use the Data.append() method.'))
+            else:
+                del self.data[index]
+        else:
+            raise TypeError(f'Found {type(index)} for index')
+
     def append(self, trial_data, strict=None):
         '''
         Append a single trial of data to the end of a Data.

tests/test_data_object.py

@@ -48,6 +48,18 @@ def test_create_outstruct_from_dict():
     assert outstruct['y'][0] == 'y0'
     assert outstruct['y'][1] == 'y1'
 
+def test_delete_field_from_outstruct():
+    data = {'x':[1,2], 'y':[3,4], 'z':[5,6]}
+    outstruct = Data(data)
+    del outstruct['z']
+    assert outstruct.fields == ['x','y']
+
+def test_delete_trial_from_outstruct():
+    data = {'x':[1,2], 'y':[3,4], 'z':[5,6]}
+    outstruct = Data(data)
+    del outstruct[0]
+    assert outstruct[0] == {'x': 2, 'y': 4, 'z': 6}
+
 def test_create_outstruct_from_dict_different_lengths():
     data = {'x': [np.array([1,2]), np.array([3,4])], 'y': ['y0', 'y1', 'y2']}
     with pytest.raises(ValueError):