Merge pull request numpy#6033 from pitrou/td64_array_repr

jaimefrio · jaimefrio · commit 61d2a445881f · 2015-07-01T14:13:36.000-07:00
ENH: improve string representation of NaTs in timedelta64 arrays
diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py
@@ -19,7 +19,8 @@
 from functools import reduce
 from . import numerictypes as _nt
 from .umath import maximum, minimum, absolute, not_equal, isnan, isinf
-from .multiarray import format_longfloat, datetime_as_string, datetime_data
+from .multiarray import (array, format_longfloat, datetime_as_string,
+                         datetime_data)
 from .fromnumeric import ravel
 from .numeric import asarray
 
@@ -733,10 +734,22 @@ def __call__(self, x):
 class TimedeltaFormat(object):
     def __init__(self, data):
         if data.dtype.kind == 'm':
-            v = data.view('i8')
-            max_str_len = max(len(str(maximum.reduce(v))),
-                              len(str(minimum.reduce(v))))
+            nat_value = array(['NaT'], dtype=data.dtype)[0]
+            v = data[not_equal(data, nat_value)].view('i8')
+            if len(v) > 0:
+                # Max str length of non-NaT elements
+                max_str_len = max(len(str(maximum.reduce(v))),
+                                  len(str(minimum.reduce(v))))
+            else:
+                max_str_len = 0
+            if len(v) < len(data):
+                # data contains a NaT
+                max_str_len = max(max_str_len, 5)
             self.format = '%' + str(max_str_len) + 'd'
+            self._nat = "'NaT'".rjust(max_str_len)
 
     def __call__(self, x):
-        return self.format % x.astype('i8')
+        if x + 1 == x:
+            return self._nat
+        else:
+            return self.format % x.astype('i8')
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
@@ -530,6 +530,17 @@ def test_datetime_array_str(self):
         a = np.array(['2010', 'NaT', '2030']).astype('M')
         assert_equal(str(a), "['2010' 'NaT' '2030']")
 
+    def test_timedelta_array_str(self):
+        a = np.array([-1, 0, 100], dtype='m')
+        assert_equal(str(a), "[ -1   0 100]")
+        a = np.array(['NaT', 'NaT'], dtype='m')
+        assert_equal(str(a), "['NaT' 'NaT']")
+        # Check right-alignment with NaTs
+        a = np.array([-1, 'NaT', 0], dtype='m')
+        assert_equal(str(a), "[   -1 'NaT'     0]")
+        a = np.array([-1, 'NaT', 1234567], dtype='m')
+        assert_equal(str(a), "[     -1   'NaT' 1234567]")
+
     def test_pickle(self):
         # Check that pickle roundtripping works
         dt = np.dtype('M8[7D]')
