Improvements to the "pretty" output in IPython

Author: apdavison

neo/core/epoch.py

@@ -6,6 +6,7 @@
 '''
 
 from copy import deepcopy, copy
+from numbers import Number
 
 import numpy as np
 import quantities as pq
@@ -88,8 +89,10 @@ def __new__(cls, times=None, durations=None, labels=None, units=None, name=None,
             raise ValueError("Times array has more than 1 dimension")
         if isinstance(durations, (list, tuple)):
             durations = np.array(durations)
-        if durations is None:
+        elif durations is None:
             durations = np.array([]) * pq.s
+        elif isinstance(durations, Number):
+            durations = durations * np.ones(times.shape)
         elif durations.size != times.size:
             if durations.size == 1:
                 durations = durations * np.ones_like(times.magnitude)
@@ -170,7 +173,14 @@ def __repr__(self):
         return '<Epoch: %s>' % ', '.join(objs)
 
     def _repr_pretty_(self, pp, cycle):
-        super()._repr_pretty_(pp, cycle)
+        labels = ""
+        if self._labels is not None:
+            labels = " with labels"
+        pp.text(f"{self.__class__.__name__} containing {self.size} epochs{labels}; "
+        f"time units {self.units.dimensionality.string}; datatype {self.dtype} ")
+        if self._has_repr_pretty_attrs_():
+            pp.breakable()
+            self._repr_pretty_attrs_(pp, cycle)
 
     def rescale(self, units):
         '''

neo/core/event.py

@@ -148,7 +148,14 @@ def __repr__(self):
         return '<Event: %s>' % ', '.join(objs)
 
     def _repr_pretty_(self, pp, cycle):
-        super()._repr_pretty_(pp, cycle)
+        labels = ""
+        if self._labels is not None:
+            labels = " with labels"
+        pp.text(f"{self.__class__.__name__} containing {self.size} events{labels}; "
+        f"time units {self.units.dimensionality.string}; datatype {self.dtype} ")
+        if self._has_repr_pretty_attrs_():
+            pp.breakable()
+            self._repr_pretty_attrs_(pp, cycle)
 
     def rescale(self, units):
         '''

neo/core/spiketrain.py

@@ -346,7 +346,21 @@ def __init__(self, times, t_stop, units=None, dtype=None, copy=True,
                             array_annotations=array_annotations, **annotations)
 
     def _repr_pretty_(self, pp, cycle):
-        super()._repr_pretty_(pp, cycle)
+        waveforms = ""
+        if self.waveforms:
+            waveforms = "with waveforms"
+        pp.text(f"{self.__class__.__name__} containing {self.size} spikes{waveforms}; "
+                f"units {self.units.dimensionality.string}; datatype {self.dtype} ")
+        if self._has_repr_pretty_attrs_():
+            pp.breakable()
+            self._repr_pretty_attrs_(pp, cycle)
+
+        def _pp(line):
+            pp.breakable()
+            with pp.group(indent=1):
+                pp.text(line)
+
+        _pp(f"time: {self.t_start} to {self.t_stop}")
 
     def rescale(self, units):
         '''