diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/io.py | 7 | ||||
| -rw-r--r-- | numpy/lib/tests/test_io.py | 16 |
2 files changed, 20 insertions, 3 deletions
diff --git a/numpy/lib/io.py b/numpy/lib/io.py index 41ffc51b8..5348866c4 100644 --- a/numpy/lib/io.py +++ b/numpy/lib/io.py @@ -312,7 +312,10 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, converters=None, """ 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 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, converters=None, # 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 diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py index 1d2c2321a..4e6412e09 100644 --- a/numpy/lib/tests/test_io.py +++ b/numpy/lib/tests/test_io.py @@ -170,7 +170,6 @@ class TestLoadTxt(TestCase): 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 @@ class TestLoadTxt(TestCase): 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 @@ class TestLoadTxt(TestCase): 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 |