[Numpy-svn] r5788 - in trunk/numpy/lib: . tests

[numpy-svn at scipy.org]

Author: dhuard
Date: 2008-09-05 08:58:00 -0500 (Fri, 05 Sep 2008)
New Revision: 5788

Modified:
   trunk/numpy/lib/io.py
   trunk/numpy/lib/tests/test_io.py
Log:
Applied patch from R. May fixing ticket #905 (loadtxt). Fixed other bug occurring when both usecols and converters are provided. Added related regression tests.

Modified: trunk/numpy/lib/io.py
===================================================================
--- trunk/numpy/lib/io.py	2008-09-05 07:00:12 UTC (rev 5787)
+++ trunk/numpy/lib/io.py	2008-09-05 13:58:00 UTC (rev 5788)
@@ -312,7 +312,10 @@
 
     """
     user_converters = converters
-
+    
+    if usecols is not None: 
+        usecols = list(usecols) 
+        
     if _string_like(fname):
         if fname.endswith('.gz'):
             import gzip
@@ -373,7 +376,7 @@
     # By preference, use the converters specified by the user
     for i, conv in (user_converters or {}).iteritems():
         if usecols:
-            i = usecols.find(i)
+            i = usecols.index(i)
         converters[i] = conv
 
     # Parse each line, including the first

Modified: trunk/numpy/lib/tests/test_io.py
===================================================================
--- trunk/numpy/lib/tests/test_io.py	2008-09-05 07:00:12 UTC (rev 5787)
+++ trunk/numpy/lib/tests/test_io.py	2008-09-05 13:58:00 UTC (rev 5788)
@@ -170,7 +170,6 @@
         a = np.array([1,2,3,4], int)
         assert_array_equal(x, a)
 
-
     def test_missing(self):
         c = StringIO.StringIO()
         c.write('1,2,3,,5\n')
@@ -180,6 +179,16 @@
         a = np.array([1,2,3,-999,5], int)
         assert_array_equal(x, a)
 
+    def test_converters_with_usecols(self):
+        c = StringIO.StringIO()
+        c.write('1,2,3,,5\n6,7,8,9,10\n')
+        c.seek(0)
+        x = np.loadtxt(c, dtype=int, delimiter=',', \
+            converters={3:lambda s: int(s or -999)}, \
+            usecols=(1, 3, ))
+        a = np.array([[2,  -999],[7, 9]], int)
+        assert_array_equal(x, a)
+        
     def test_comments(self):
         c = StringIO.StringIO()
         c.write('# comment\n1,2,3,5\n')
@@ -221,6 +230,11 @@
         x = np.loadtxt(c, dtype=float, usecols=(1,2))
         assert_array_equal(x, a[:,1:])
 
+        # Testing with arrays instead of tuples. 
+        c.seek(0)
+        x = np.loadtxt(c, dtype=float, usecols=np.array([1,2]))
+        assert_array_equal(x, a[:,1:])
+
         # Checking with dtypes defined converters.
         data = '''JOE 70.1 25.3
                 BOB 60.5 27.9