Skip to content

Commit b3e2438

Browse files
committed
Accept range for list-requiring args in py3
1 parent 13b57cd commit b3e2438

File tree

2 files changed

+22
-2
lines changed

2 files changed

+22
-2
lines changed

pandas/io/parsers.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -943,6 +943,8 @@ def _clean_options(self, options, engine):
943943
if _is_index_col(index_col):
944944
if not isinstance(index_col, (list, tuple, np.ndarray)):
945945
index_col = [index_col]
946+
if PY3 and isinstance(index_col, range):
947+
index_col = list(index_col)
946948
result['index_col'] = index_col
947949

948950
names = list(names) if names is not None else names
@@ -1191,6 +1193,8 @@ def __init__(self, kwds):
11911193

11921194
# validate header options for mi
11931195
self.header = kwds.get('header')
1196+
if PY3 and isinstance(self.header, range):
1197+
self.header = list(self.header)
11941198
if isinstance(self.header, (list, tuple, np.ndarray)):
11951199
if not all(map(is_integer, self.header)):
11961200
raise ValueError("header must be integer or list of integers")
@@ -1213,7 +1217,6 @@ def __init__(self, kwds):
12131217
is_integer(self.index_col)):
12141218
raise ValueError("index_col must only contain row numbers "
12151219
"when specifying a multi-index header")
1216-
12171220
# GH 16338
12181221
elif self.header is not None and not is_integer(self.header):
12191222
raise ValueError("header must be integer or list of integers")

pandas/tests/io/parser/index_col.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
import pandas.util.testing as tm
1212

13-
from pandas import DataFrame, Index, MultiIndex
13+
from pandas import DataFrame, Index, MultiIndex, date_range
1414
from pandas.compat import StringIO
1515

1616

@@ -141,3 +141,20 @@ def test_empty_with_index_col_false(self):
141141
result = self.read_csv(StringIO(data), index_col=False)
142142
expected = DataFrame([], columns=['x', 'y'])
143143
tm.assert_frame_equal(result, expected)
144+
145+
def test_range_index_col(self):
146+
cols = MultiIndex.from_arrays([['A', 'B', 'C'], ['foo', 'bar', 'baz']])
147+
index = date_range('2016-01-02', periods=3, freq='D')
148+
data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
149+
df = DataFrame(data, index=index, columns=cols)
150+
151+
data = (",A,B,C\n"
152+
",foo,bar,baz\n"
153+
"2016-01-02,1,2,3\n"
154+
"2016-01-03,4,5,6\n"
155+
"2016-01-04,7,8,9"
156+
)
157+
res = self.read_csv(StringIO(df.to_csv()),
158+
index_col=range(1),
159+
header=range(2))
160+
tm.assert_frame_equal(df, res)

0 commit comments

Comments
 (0)