Merge pull request #269 from ndawe/master

stretch: if fields is single string (not list) then flatten output

root_numpy/_utils.py

@@ -126,8 +126,11 @@ def stretch(arr, fields=None, return_indices=False):
     ----------
     arr : NumPy structured or record array
         The array to be stretched.
-    fields : list of strings, optional (default=None)
-        A list of column names to stretch. If None, then stretch all fields.
+    fields : list of strings or single string, optional (default=None)
+        A list of column names or a single column name to stretch.
+        If ``fields`` is a string, then the output array is a one-dimensional
+        unstructured array containing only the stretched elements of that
+        field. If None, then stretch all fields.
     return_indices : bool, optional (default=False)
         If True, the array index of each stretched array entry will be
         returned in addition to the stretched array.
@@ -154,8 +157,12 @@ def stretch(arr, fields=None, return_indices=False):
     dtype = []
     len_array = None
 
+    flatten = False
     if fields is None:
         fields = arr.dtype.names
+    elif isinstance(fields, string_types):
+        fields = [fields]
+        flatten = True
 
     # Construct dtype and check consistency
     for field in fields:
@@ -197,6 +204,9 @@ def stretch(arr, fields=None, return_indices=False):
             # Scalar field
             ret[field] = np.repeat(arr[field], len_array)
 
+    if flatten:
+        ret = ret[fields[0]]
+
     if return_indices:
         idx = np.concatenate(list(map(np.arange, len_array)))
         return ret, idx

root_numpy/tests.py

@@ -704,6 +704,10 @@ def test_stretch():
     from_stretched = stretched[idx == 0]['vl1']
     assert_array_equal(from_arr, from_stretched)
 
+    # stretch single field and produce unstructured output
+    stretched = rnp.stretch(arr, 'vl1')
+    assert_equal(stretched.dtype, np.int)
+
 
 def test_blockwise_inner_join():
     test_data = np.array([