Add more dict methods to datagroup and dataset

src/osyris/core/datagroup.py

@@ -64,6 +64,14 @@ def __str__(self):
             output += str(item) + "\n"
         return output
 
+    def __eq__(self, other):
+        if self.keys() != other.keys():
+            return False
+        for key, value in self.items():
+            if all(value != other[key]):
+                return False
+        return True
+
     @property
     def parent(self):
         return self._parent
@@ -106,3 +114,17 @@ def sortby(self, key):
                 key = np.argsort(self[key]).values
             for var in self.keys():
                 self[var] = self[var][key]
+
+    def clear(self):
+        self._container.clear()
+        self.shape = None
+
+    def get(self, key, default):
+        return self._container.get(key, default)
+
+    def pop(self, key):
+        return self._container.pop(key)
+
+    def update(self, d):
+        for key, value in d.items():
+            self[key] = value

src/osyris/core/dataset.py

@@ -71,3 +71,17 @@ def nbytes(self):
 
     def print_size(self):
         return bytes_to_human_readable(self.nbytes())
+
+    def clear(self):
+        self.groups.clear()
+        self.meta.clear()
+
+    def get(self, key, default):
+        return self.groups.get(key, default)
+
+    def pop(self, key):
+        return self.groups.pop(key)
+
+    def update(self, d):
+        for key, value in d.items():
+            self[key] = value

test/test_datagroup.py

@@ -33,9 +33,7 @@ def test_datagroup_insert_vector():
     b = Array(values=[[6., 7., 8.], [9., 10., 11.], [12., 13., 14.], [15., 16., 17.],
                       [18., 19., 20.]],
               unit='s')
-    dg = Datagroup()
-    dg['a'] = a
-    dg['b'] = b
+    dg = Datagroup({'a': a, 'b': b})
     assert 'a' in dg.keys()
     assert 'b' in dg.keys()
     assert dg['a'].name == 'a'
@@ -55,43 +53,39 @@ def test_datagroup_bad_length_insertion():
 def test_datagroup_delitem():
     a = Array(values=[1., 2., 3., 4., 5.], unit='m')
     b = Array(values=[6., 7., 8., 9., 10.], unit='s')
-    dg = Datagroup()
-    dg['a'] = a
-    dg['b'] = b
+    dg = Datagroup({'a': a, 'b': b})
     del dg['a']
     assert 'a' not in dg.keys()
     assert 'b' in dg.keys()
     assert len(dg) == 1
 
 
+def test_datagroup_equal_operator():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    assert Datagroup({'a': a, 'b': b}) == Datagroup({'a': a, 'b': b})
+
+
 def test_datagroup_slice():
     a = Array(values=[1., 2., 3., 4., 5.], unit='m')
     b = Array(values=[6., 7., 8., 9., 10.], unit='s')
-    dg = Datagroup()
-    dg['a'] = a
-    dg['b'] = b
+    dg = Datagroup({'a': a, 'b': b})
     expected = Datagroup({
         'a': Array(values=[2.], unit='m'),
         'b': Array(values=[7.], unit='s')
     })
-    sliced = dg[1]
-    assert sliced['a'] == expected['a']
-    assert sliced['b'] == expected['b']
+    assert dg[1] == expected
 
 
 def test_datagroup_slice_range():
     a = Array(values=[1., 2., 3., 4., 5.], unit='m')
     b = Array(values=[6., 7., 8., 9., 10.], unit='s')
-    dg = Datagroup()
-    dg['a'] = a
-    dg['b'] = b
+    dg = Datagroup({'a': a, 'b': b})
     expected = Datagroup({
         'a': Array(values=[2., 3., 4.], unit='m'),
         'b': Array(values=[7., 8., 9.], unit='s')
     })
-    sliced = dg[1:4]
-    assert all(sliced['a'] == expected['a'])
-    assert all(sliced['b'] == expected['b'])
+    assert dg[1:4] == expected
 
 
 def test_datagroup_sortby_key():
@@ -118,3 +112,40 @@ def test_datagroup_sortby_indices():
     })
     assert all(dg['a'] == expected['a'])
     assert all(dg['b'] == expected['b'])
+
+
+def test_datagroup_clear():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    dg = Datagroup({'a': a, 'b': b})
+    dg.clear()
+    assert len(dg) == 0
+
+
+def test_datagroup_get():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    c = Array(values=[11., 12., 13., 14., 15.], unit='K')
+    dg = Datagroup({'a': a, 'b': b})
+    assert all(dg.get('a', 42) == a)
+    assert dg.get('c', 42) == 42
+    assert all(dg.get('c', c) == c)
+
+
+def test_datagroup_pop():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    dg = Datagroup({'a': a, 'b': b})
+    c = dg.pop('a')
+    assert 'a' not in dg
+    assert all(c == a)
+
+
+def test_datagroup_update():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    c = Array(values=[11., 12., 13., 14., 15.], unit='K')
+    dg = Datagroup({'a': a, 'b': b})
+    dg.update({'b': 2.0 * b, 'c': c})
+    assert all(dg['b'] == Array(values=[12., 14., 16., 18., 20.], unit='s'))
+    assert all(dg['c'] == c)

test/test_dataset.py

@@ -1,19 +1,68 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2022 Osyris contributors (https://github.com/osyris-project/osyris)
-import osyris
+from osyris import Array, Datagroup, Dataset
 
 
 def test_dataset_creation():
-    a = osyris.Array(values=[1., 2., 3., 4., 5.], unit='m')
-    b = osyris.Array(values=[6., 7., 8., 9., 10.], unit='s')
-    dg1 = osyris.Datagroup()
-    dg1['a'] = a
-    dg1['b'] = b
-    dg2 = osyris.Datagroup()
-    dg2['a'] = a[:-1]
-    dg2['b'] = b[:-1]
-    ds = osyris.Dataset()
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    dg1 = Datagroup({'a': a, 'b': b})
+    dg2 = Datagroup({'a': a[:-1], 'b': b[:-1]})
+    ds = Dataset()
     ds['one'] = dg1
     ds['two'] = dg2
     assert 'one' in ds
     assert 'two' in ds
+
+
+def test_dataset_clear():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    dg1 = Datagroup({'a': a, 'b': b})
+    dg2 = Datagroup({'a': a[:-1], 'b': b[:-1]})
+    ds = Dataset()
+    ds['one'] = dg1
+    ds['two'] = dg2
+    ds.clear()
+    assert len(ds) == 0
+
+
+def test_dataset_get():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    dg1 = Datagroup({'a': a, 'b': b})
+    dg2 = Datagroup({'a': a[:-1], 'b': b[:-1]})
+    ds = Dataset()
+    ds['one'] = dg1
+    ds['two'] = dg2
+    assert ds.get('one', 42) == dg1
+    assert ds.get('three', 42) == 42
+
+
+def test_dataset_pop():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    dg1 = Datagroup({'a': a, 'b': b})
+    dg2 = Datagroup({'a': a[:-1], 'b': b[:-1]})
+    ds = Dataset()
+    ds['one'] = dg1
+    ds['two'] = dg2
+    c = ds.pop('one')
+    assert 'one' not in ds
+    assert c == dg1
+
+
+def test_dataset_update():
+    a = Array(values=[1., 2., 3., 4., 5.], unit='m')
+    b = Array(values=[6., 7., 8., 9., 10.], unit='s')
+    c = Array(values=[11., 12., 13., 14., 15.], unit='K')
+    dg1 = Datagroup({'a': a, 'b': b})
+    dg2 = Datagroup({'a': a[:-1], 'b': b[:-1]})
+    dg3 = Datagroup({'a': a, 'c': c})
+    ds = Dataset()
+    ds['one'] = dg1
+    ds['two'] = dg2
+    ds.update({'one': Datagroup({'a': a * 2.0}), 'three': dg3})
+    assert 'three' in ds
+    assert 'b' not in ds['one']
+    assert 'c' in ds['three']