diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/lib/npyio.py | 35 | ||||
| -rw-r--r-- | numpy/lib/tests/test_io.py | 38 |
2 files changed, 68 insertions, 5 deletions
diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py index 640f4fa32..4db542b55 100644 --- a/numpy/lib/npyio.py +++ b/numpy/lib/npyio.py @@ -6,7 +6,7 @@ import re import itertools import warnings import weakref -from operator import itemgetter +from operator import itemgetter, index as opindex import numpy as np from . import format @@ -714,10 +714,18 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, ``converters = {3: lambda s: float(s.strip() or 0)}``. Default: None. skiprows : int, optional Skip the first `skiprows` lines; default: 0. - usecols : sequence, optional - Which columns to read, with 0 being the first. For example, - ``usecols = (1,4,5)`` will extract the 2nd, 5th and 6th columns. + + usecols : int or sequence, optional + Which columns to read, with 0 being the first. For example, + usecols = (1,4,5) will extract the 2nd, 5th and 6th columns. The default, None, results in all columns being read. + + .. versionadded:: 1.11.0 + + Also when a single column has to be read it is possible to use + an integer instead of a tuple. E.g ``usecols = 3`` reads the + third column the same way as `usecols = (3,)`` would. + unpack : bool, optional If True, the returned array is transposed, so that arguments may be unpacked using ``x, y, z = loadtxt(...)``. When used with a structured @@ -786,8 +794,25 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, user_converters = converters if delimiter is not None: delimiter = asbytes(delimiter) + if usecols is not None: - usecols = list(usecols) + # Allow usecols to be a single int or a sequence of ints + try: + usecols_as_list = list(usecols) + except TypeError: + usecols_as_list = [usecols] + for col_idx in usecols_as_list: + try: + opindex(col_idx) + except TypeError as e: + e.args = ( + "usecols must be an int or a sequence of ints but " + "it contains at least one element of type %s" % + type(col_idx), + ) + raise + # Fall back to existing code + usecols = usecols_as_list fown = False try: diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py index 226dc88fa..c0f8c1953 100644 --- a/numpy/lib/tests/test_io.py +++ b/numpy/lib/tests/test_io.py @@ -608,6 +608,29 @@ class TestLoadTxt(TestCase): x = np.loadtxt(c, dtype=float, usecols=np.array([1, 2])) assert_array_equal(x, a[:, 1:]) + # Testing with an integer instead of a sequence + for int_type in [int, np.int8, np.int16, + np.int32, np.int64, np.uint8, np.uint16, + np.uint32, np.uint64]: + to_read = int_type(1) + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=to_read) + assert_array_equal(x, a[:, 1]) + + # Testing with some crazy custom integer type + class CrazyInt(object): + def __index__(self): + return 1 + + crazy_int = CrazyInt() + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=crazy_int) + assert_array_equal(x, a[:, 1]) + + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=(crazy_int,)) + assert_array_equal(x, a[:, 1]) + # Checking with dtypes defined converters. data = '''JOE 70.1 25.3 BOB 60.5 27.9 @@ -619,6 +642,21 @@ class TestLoadTxt(TestCase): assert_equal(arr['stid'], [b"JOE", b"BOB"]) assert_equal(arr['temp'], [25.3, 27.9]) + # Testing non-ints in usecols + c.seek(0) + bogus_idx = 1.5 + assert_raises_regex( + TypeError, + '^usecols must be.*%s' % type(bogus_idx), + np.loadtxt, c, usecols=bogus_idx + ) + + assert_raises_regex( + TypeError, + '^usecols must be.*%s' % type(bogus_idx), + np.loadtxt, c, usecols=[0, bogus_idx, 0] + ) + def test_fancy_dtype(self): c = TextIO() c.write('1,2,3.0\n4,5,6.0\n') |