Repository: l-johnston/nitdms

Issue:
Add head() and tail() functions to WaveformDT. AB#5

Discussion:

Patch:
diff --git a/nitdms/waveformdatatype.py b/nitdms/waveformdatatype.py
index 75f246e..7dbc4c1 100644
--- a/nitdms/waveformdatatype.py
+++ b/nitdms/waveformdatatype.py
@@ -46,9 +46,24 @@ class WaveformDT(np.ndarray):
         self.dt = getattr(obj, "dt", 1.0)
 
     def __repr__(self):
-        repr_str = super(WaveformDT, self).__repr__()
-        wf_details = f", {self.dt}, {self.t0})"
-        return repr_str.replace(")", wf_details)
+        t0 = self.t0 if isinstance(self.t0, Number) else 0.0
+        t0 += getattr(self, "wf_start_offset", 0.0)
+        dt = self.dt
+        rows = []
+        if self.size < 50:
+            for i, sample in enumerate(self):
+                rows.append(f"{t0 + i*dt:11.4e}\t{sample:11.4e}")
+        else:
+            for i, sample in enumerate(self[:5]):
+                rows.append(f"{t0 + i*dt:11.4e}\t{sample:11.4e}")
+            rows.append(" ...")
+            t0 = t0 + (self.size - 5) * dt
+            for i, sample in enumerate(self[-5:]):
+                rows.append(f"{t0 + i*dt:11.4e}\t{sample:11.4e}")
+        rows.append(f"Length: {self.size}")
+        rows.append(f"t0: {self.t0}")
+        rows.append(f"dt: {self.dt:11.4e}")
+        return "\n".join(rows)
 
     @property
     def Y(self):
@@ -122,3 +137,30 @@ class WaveformDT(np.ndarray):
         if method == "at":
             return None
         return results[0] if len(results) == 1 else results
+
+    def head(self, n=5):
+        """Return first n samples of the waveform
+
+        Args:
+            n (int): number of samples to return
+
+        Returns:
+            (WaveformDT): first n samples
+        """
+        return self[:n]
+
+    def tail(self, n=5):
+        """Return the last n samples of the waveform
+
+        Args:
+            n (int): number of samples to return
+
+        Returns:
+            (WaveformDT): last n samples
+        """
+        start_offset = self.t0 if isinstance(self.t0, Number) else 0.0
+        start_offset += getattr(self, "wf_start_offset", 0.0)
+        start_offset += (self.size - n) * self.dt
+        wf = self[-n:]
+        setattr(wf, "wf_start_offset", start_offset)
+        return wf


Test patch:
diff --git a/tests/test_wdt.py b/tests/test_wdt.py
index 030f451..3de2f01 100644
--- a/tests/test_wdt.py
+++ b/tests/test_wdt.py
@@ -65,8 +65,28 @@ def test_extraattributes():
 
 
 def test_repr():
-    wf = WaveformDT([1, 2, 3], 1, 0)
-    assert wf.__repr__() == "WaveformDT([1, 2, 3], 1, 0)"
+    wf = WaveformDT([1, 2], 1, 0)
+    expected = (
+        " 0.0000e+00\t 1.0000e+00\n 1.0000e+00\t 2.0000e+00\n"
+        + "Length: 2\nt0: 0\ndt:  1.0000e+00"
+    )
+    assert wf.__repr__() == expected
+    wf = WaveformDT(range(60), 1, 0)
+    expected = (
+        " 0.0000e+00\t 0.0000e+00\n"
+        " 1.0000e+00\t 1.0000e+00\n"
+        " 2.0000e+00\t 2.0000e+00\n"
+        " 3.0000e+00\t 3.0000e+00\n"
+        " 4.0000e+00\t 4.0000e+00\n"
+        " ...\n"
+        " 5.5000e+01\t 5.5000e+01\n"
+        " 5.6000e+01\t 5.6000e+01\n"
+        " 5.7000e+01\t 5.7000e+01\n"
+        " 5.8000e+01\t 5.8000e+01\n"
+        " 5.9000e+01\t 5.9000e+01\n"
+        "Length: 60\nt0: 0\ndt:  1.0000e+00"
+    )
+    assert wf.__repr__() == expected
 
 
 def test_Y():
@@ -103,3 +123,31 @@ def test_ufunc_out():
     np.multiply(a, 2.0, out=c)
     results = c == WaveformDT([2.0, 4.0, 6.0], 2, 1)
     assert results.all()
+
+
+def test_head():
+    wf = WaveformDT(range(60), 1, 0)
+    wf_head = wf.head()
+    expected = (
+        " 0.0000e+00\t 0.0000e+00\n"
+        " 1.0000e+00\t 1.0000e+00\n"
+        " 2.0000e+00\t 2.0000e+00\n"
+        " 3.0000e+00\t 3.0000e+00\n"
+        " 4.0000e+00\t 4.0000e+00\n"
+        "Length: 5\nt0: 0\ndt:  1.0000e+00"
+    )
+    assert wf_head.__repr__() == expected
+
+
+def test_tail():
+    wf = WaveformDT(range(60), 1, 0)
+    wf_tail = wf.tail()
+    expected = (
+        " 5.5000e+01\t 5.5000e+01\n"
+        " 5.6000e+01\t 5.6000e+01\n"
+        " 5.7000e+01\t 5.7000e+01\n"
+        " 5.8000e+01\t 5.8000e+01\n"
+        " 5.9000e+01\t 5.9000e+01\n"
+        "Length: 5\nt0: 0\ndt:  1.0000e+00"
+    )
+    assert wf_tail.__repr__() == expected