Merge pull request #92 from bmcfee/pumpp-repr

Pumpp repr

pumpp/core.py

@@ -11,6 +11,7 @@
 
 import librosa
 import jams
+import six
 
 from .base import Slicer
 from .exceptions import ParameterError
@@ -225,3 +226,28 @@ def __getitem__(self, key):
 
     def __call__(self, *args, **kwargs):
         return self.transform(*args, **kwargs)
+
+    def __str__(self):
+        rstr = '<Pump [{:d} operators, {:d} fields]>'.format(len(self.ops),
+                                                             len(self.fields))
+        for key in self.opmap:
+            rstr += "\n  - '{}': {}".format(key, type(self.opmap[key]))
+            for field in self.opmap[key].fields:
+                rstr += "\n    - '{}': {}".format(field, self.opmap[key].fields[field])
+        return rstr
+
+    def _repr_html_(self):
+
+        rstr = '<dl class="row">'
+        for key in self.opmap:
+            rstr += '\n  <dt class="col-sm-3">{:s}</dt>'.format(key)
+            rstr += '\n  <dd class="col-sm-9">{}'.format(self.opmap[key])
+
+            rstr += '<ul>'
+            for fkey, field in six.iteritems(self.opmap[key].fields):
+                rstr += '\n  <li>{:s} [shape={}, dtype={}]</li>'.format(fkey,
+                                                                        field.shape,
+                                                                        field.dtype.__name__)
+            rstr += '</ul></dd>'
+        rstr += '</dl>'
+        return rstr

tests/test_core.py

@@ -212,3 +212,45 @@ def test_pump_layers(sr, hop_length):
         assert L1[k].dtype == L2[k].dtype
         for d1, d2 in zip(L1[k].shape, L2[k].shape):
             assert str(d1) == str(d2)
+
+
+def test_pump_str(sr, hop_length):
+
+    ops = [pumpp.feature.STFT(name='stft', sr=sr,
+                              hop_length=hop_length,
+                              n_fft=2*hop_length),
+
+           pumpp.task.BeatTransformer(name='beat', sr=sr,
+                                      hop_length=hop_length),
+
+           pumpp.task.ChordTransformer(name='chord', sr=sr,
+                                       hop_length=hop_length),
+
+           pumpp.task.StaticLabelTransformer(name='tags',
+                                             namespace='tag_open',
+                                             labels=['rock', 'jazz'])]
+
+    pump = pumpp.Pump(*ops)
+
+    assert isinstance(str(pump), str)
+
+
+def test_pump_repr_html(sr, hop_length):
+
+    ops = [pumpp.feature.STFT(name='stft', sr=sr,
+                              hop_length=hop_length,
+                              n_fft=2*hop_length),
+
+           pumpp.task.BeatTransformer(name='beat', sr=sr,
+                                      hop_length=hop_length),
+
+           pumpp.task.ChordTransformer(name='chord', sr=sr,
+                                       hop_length=hop_length),
+
+           pumpp.task.StaticLabelTransformer(name='tags',
+                                             namespace='tag_open',
+                                             labels=['rock', 'jazz'])]
+
+    pump = pumpp.Pump(*ops)
+
+    assert isinstance(pump._repr_html_(), str)