Merge pull request numpy#6643 from ahaldane/recarray_getitem_returns_recarray

charris · charris · commit 9cc7bc3ad4e3 · 2015-11-06T11:21:22.000-07:00
ENH: make recarray.getitem return a recarray
diff --git a/numpy/core/records.py b/numpy/core/records.py
@@ -502,6 +502,7 @@ def __getitem__(self, indx):
         # we might also be returning a single element
         if isinstance(obj, ndarray):
             if obj.dtype.fields:
+                obj = obj.view(recarray)
                 if issubclass(obj.dtype.type, nt.void):
                     return obj.view(dtype=(self.dtype.type, obj.dtype))
                 return obj
diff --git a/numpy/core/tests/test_records.py b/numpy/core/tests/test_records.py
@@ -121,6 +121,14 @@ def test_recarray_views(self):
         assert_equal(type(rv), np.recarray)
         assert_equal(rv.dtype.type, np.record)
 
+        #check that getitem also preserves np.recarray and np.record
+        r = np.rec.array(np.ones(4, dtype=[('a', 'i4'), ('b', 'i4'), 
+                                           ('c', 'i4,i4')]))
+        assert_equal(r['c'].dtype.type, np.record)
+        assert_equal(type(r['c']), np.recarray)
+        assert_equal(r[['a', 'b']].dtype.type, np.record)
+        assert_equal(type(r[['a', 'b']]), np.recarray)
+
         # check that accessing nested structures keep record type, but
         # not for subarrays, non-void structures, non-structured voids
         test_dtype = [('a', 'f4,f4'), ('b', 'V8'), ('c', ('f4',2)),
