SF.net SVN: matplotlib: [5212] trunk/matplotlib/lib/matplotlib/mlab.py

[<sa...@us...>]

Revision: 5212
          http://matplotlib.svn.sourceforge.net/matplotlib/?rev=5212&view=rev
Author:   sameerd
Date:     2008-05-21 09:08:26 -0700 (Wed, 21 May 2008)

Log Message:
-----------
Record array append multiple fields 

Modified Paths:
--------------
    trunk/matplotlib/lib/matplotlib/mlab.py

Modified: trunk/matplotlib/lib/matplotlib/mlab.py
===================================================================
--- trunk/matplotlib/lib/matplotlib/mlab.py	2008-05-21 13:06:51 UTC (rev 5211)
+++ trunk/matplotlib/lib/matplotlib/mlab.py	2008-05-21 16:08:26 UTC (rev 5212)
@@ -1953,14 +1953,39 @@
 
 def rec_append_field(rec, name, arr, dtype=None):
     'return a new record array with field name populated with data from array arr'
-    arr = np.asarray(arr)
-    if dtype is None:
-        dtype = arr.dtype
-    newdtype = np.dtype(rec.dtype.descr + [(name, dtype)])
+    warnings.warn("use rec_append_fields", DeprecationWarning)
+    return rec_append_fields(rec, name, arr, dtype)
+
+def rec_append_fields(rec, names, arrs, dtypes=None):
+    """
+    return a new record array with field names populated with data
+    from arrays in arrs.  If appending a single field then names, arrs
+    and dtypes do not have to be lists. They can just be the values themselves.
+    """
+    if (not cbook.is_string_like(names) and cbook.iterable(names) \
+            and len(names) and cbook.is_string_like(names[0])):
+        if len(names) != len(arrs):
+            raise ValueError, "number of arrays do not match number of names"
+    else: # we have only 1 name and 1 array
+        names = [names]
+        arrs = [arrs]
+    arrs = map(np.asarray, arrs)
+    if dtypes is None:
+        dtypes = [a.dtype for a in arrs]
+    elif not cbook.iterable(dtypes):
+        dtypes = [dtypes]
+    if len(arrs) != len(dtypes):
+        if len(dtypes) == 1:
+            dtypes = dtypes * len(arrs)
+        else:
+            raise ValueError, "dtypes must be None, a single dtype or a list"
+
+    newdtype = np.dtype(rec.dtype.descr + zip(names, dtypes))
     newrec = np.empty(rec.shape, dtype=newdtype)
     for field in rec.dtype.fields:
         newrec[field] = rec[field]
-    newrec[name] = arr
+    for name, arr in zip(names, arrs):
+        newrec[name] = arr
     return newrec.view(np.recarray)
 
 


This was sent by the SourceForge.net collaborative development platform, the world's largest Open Source development site.