From 96d1f9afd7e06f343bdb335aea78301268fc1513 Mon Sep 17 00:00:00 2001
From: wirawan <wirawan>
Date: Thu, 30 Sep 2010 17:23:34 +0000
Subject: [PATCH] * Upgrade to allow reading complex-valued data directly.  
 Please see the documentation of text_input.read_items for more detail.

---
 iofmt/text_input.py | 61 ++++++++++++++++++++++++++++++++++-----------
 1 file changed, 46 insertions(+), 15 deletions(-)

diff --git a/iofmt/text_input.py b/iofmt/text_input.py
index ac4ff29..99bc439 100644
--- a/iofmt/text_input.py
+++ b/iofmt/text_input.py
@@ -1,5 +1,5 @@
 #!/usr/bin/python
-# $Id: text_input.py,v 1.1 2010-09-27 19:54:05 wirawan Exp $
+# $Id: text_input.py,v 1.2 2010-09-30 17:23:34 wirawan Exp $
 #
 # wpylib.iofmt.text_input module
 # Quick-n-dirty text input utilities
@@ -146,6 +146,13 @@ class text_input(object):
     If the tuple contains the third field, it is used as the name of the field;
     otherwise the fields are named f0, f1, f2, ....
 
+    Complex data (floating-point only) must be specified as a tuple of two columns
+    containing the real and imaginary data, like this:
+       ((2, 3), complex, 'ampl')
+    or
+       ((7, 9), complex)     # fine to interleave column with something else
+
+
     Additional keyword options:
     * deftype: default datatype
     * maxcount: maximum number of records to be read
@@ -154,26 +161,50 @@ class text_input(object):
     """
     deftype = kwd.get("deftype", float)
 
-    # float_fields extracts the desired columns and converts them to floats
-    flds = []
-    cols = []
+    class register_item_t:
+      flds = []
+      cols = []
+      complex_types = (complex, numpy.complexfloating)
+      def add(self, col, fldname, type):
+        dtype = numpy.dtype(type)
+        t = dtype.type
+        dsamp = t() # create a sample
+        # Special handling for complex:
+        # -- unfortunately this detection fails because even real
+        # numbers have its 'imag' attribute:
+        #dattrs = dir(dsamp)
+        #if "imag" in dattrs and "real" in dattrs:
+        if isinstance(dsamp, numpy.complexfloating):
+          dtype_elem = dsamp.real.dtype
+          t_elem = dtype_elem.type
+          conv_func = lambda v, c: t(t_elem(v[c[0]]) + 1j*t_elem(v[c[1]]))
+          self.cols.append((conv_func, col))
+          self.flds.append((fldname, dtype))
+        else:
+          # other datatypes: much easier
+          # Simply get the string, and use numpy to convert to the datatype
+          # on-the-fly
+          conv_func = lambda v, c: t(v[c])
+          self.cols.append((conv_func, col))
+          self.flds.append((fldname, dtype))
+    reg = register_item_t()
+
     for (i,c) in zip(xrange(len(col_desc)), col_desc):
       if type(c) == int:
-        cols.append(c)
-        flds.append(('f' + str(i), deftype))
+        reg.add(c, 'f' + str(i), deftype)
       elif len(c) == 1:
-        cols.append(c[0])
-        flds.append(('f' + str(i), deftype))
+        reg.add(c[0], 'f' + str(i), deftype)
       elif len(c) == 2:
-        cols.append(c[0])
-        flds.append(('f' + str(i), c[1]))
+        reg.add(c[0], 'f' + str(i), c[1])
       elif len(c) == 3:
-        cols.append(c[0])
-        flds.append((c[2], c[1]))
+        reg.add(c[0], c[2], c[1])
+      else:
+        raise ValueError, \
+          "Invalid column specification: %s" % (c,)
 
-    #print cols
-    #print flds
-    get_fields = lambda vals : tuple([ vals[col] for col in cols ])
+    cols = reg.cols
+    flds = reg.flds
+    get_fields = lambda vals : tuple([ filt(vals,col) for (filt,col) in cols ])
     if "maxcount" in kwd:
       #print "hello"
       rslt = [ get_fields(vals.split()) for (c,vals) in zip(xrange(kwd['maxcount']),self) ]