Change __getitem__ to handle a single item or a slice, including mult…

…i-channel case. Apply suggestions from code review and add test. Co-authored-by: Teque5 <[email protected]>
sigmf · Oct 25, 2023 · 6ee4909 · 6ee4909
1 parent 406eed7
commit 6ee4909
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 12 deletions.
diff --git a/sigmf/sigmffile.py b/sigmf/sigmffile.py
@@ -201,18 +201,21 @@ def __next__(self):
             raise StopIteration
 
     def __getitem__(self, sli):
-        a = self._memmap[sli] # matches behavior of numpy.ndarray.__getitem__()
-        if self._return_type is not None:
-            # is_fixed_point and is_complex
-            if self._memmap.ndim == 2:
-                # num_channels == 1
-                a = a[:,0].astype(self._return_type) + 1.j * a[:,1].astype(self._return_type)
-            elif self._memmap.ndim == 3:
-                # num_channels > 1
-                a = a[:,:,0].astype(self._return_type) + 1.j * a[:,:,1].astype(self._return_type)
-            else:
-                raise ValueError("unhandled ndim in SigMFFile.__getitem__(); this shouldn't happen")
-        return a
+        mem = self._memmap[sli] # matches behavior of numpy.ndarray.__getitem__()
+
+        if self._return_type is None:
+            return mem
+
+        # is_fixed_point and is_complex
+        if self._memmap.ndim == 2:
+            # num_channels == 1
+            ray = mem[:,0].astype(self._return_type) + 1.j * mem[:,1].astype(self._return_type)
+        elif self._memmap.ndim == 3:
+            # num_channels > 1
+            ray = mem[:,:,0].astype(self._return_type) + 1.j * mem[:,:,1].astype(self._return_type)
+        else:
+            raise ValueError("unhandled ndim in SigMFFile.__getitem__(); this shouldn't happen")
+        return ray[0] if type(sli) is int else ray # return element instead of 1-element array
 
     def _get_start_offset(self):
         """

diff --git a/tests/test_sigmffile.py b/tests/test_sigmffile.py
@@ -246,3 +246,26 @@ def test_captures_checking():
     assert (160,224) == sigmf4.get_capture_byte_boundarys(1)
     assert np.array_equal(np.array(range(64)), sigmf4.read_samples_in_capture(0,autoscale=False)[:,0])
     assert np.array_equal(np.array(range(64,96)), sigmf4.read_samples_in_capture(1,autoscale=False)[:,1])
+
+def test_slicing():
+    '''Test __getitem___ builtin for sigmffile, make sure slicing and indexing works as expected.'''
+    np.array(TEST_U8_DATA0, dtype=np.uint8).tofile('/tmp/d5.sigmf-data')
+    with open('/tmp/d5.sigmf-meta','w') as f0: json.dump(TEST_U8_META0, f0)
+    sigmf0 = sigmffile.fromfile('/tmp/d5.sigmf-meta')
+    assert np.array_equal(TEST_U8_DATA0, sigmf0[:])
+    assert TEST_U8_DATA0[6] == sigmf0[6]
+
+    np.array(TEST_FLOAT32_DATA, dtype=np.float32).tofile('/tmp/d6.sigmf-data')
+    with open('/tmp/d6.sigmf-meta','w') as f1: json.dump(TEST_METADATA, f1)
+    sigmf1 = sigmffile.fromfile('/tmp/d6.sigmf-meta')
+    assert np.array_equal(TEST_FLOAT32_DATA, sigmf1[:])
+    assert sigmf1[10] == TEST_FLOAT32_DATA[10]
+
+    np.array(TEST_U8_DATA4, dtype=np.uint8).tofile('/tmp/d7.sigmf-data')
+    with open('/tmp/d7.sigmf-meta','w') as f4: json.dump(TEST_U8_META4, f4)
+    sigmf2 = sigmffile.fromfile('/tmp/d7.sigmf-meta')
+    channelized = np.array(TEST_U8_DATA4).reshape((128,2))
+    assert np.array_equal(channelized, sigmf2[:][:])
+    assert np.array_equal(sigmf2[10:20, 91:112], sigmf2.read_samples(autoscale=False)[10:20, 91:112])
+    assert np.array_equal(sigmf2[0], channelized[0])
+    assert np.array_equal(sigmf2[1,:], channelized[1,:])