From 710a7a83f3bfec0b8a0c7680d1b3466ecee7a987 Mon Sep 17 00:00:00 2001 From: Mo Ip Date: Mon, 22 May 2017 10:58:49 -0700 Subject: [PATCH 1/2] sets data index if it's not None --- pandas/io/stata.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pandas/io/stata.py b/pandas/io/stata.py index e03e87f09173e..4916da375f0f4 100644 --- a/pandas/io/stata.py +++ b/pandas/io/stata.py @@ -178,6 +178,8 @@ def read_stata(filepath_or_buffer, convert_dates=True, data = reader.read() finally: reader.close() + if index is not None: + data.set_index(index) return data From 87a2de361e1165dfc9711a6cadadd93e89711c3a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 May 2017 13:14:34 -0500 Subject: [PATCH 2/2] TST: add basic test --- pandas/tests/io/test_stata.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pandas/tests/io/test_stata.py b/pandas/tests/io/test_stata.py index b9c6736563160..24dd9a6700482 100644 --- a/pandas/tests/io/test_stata.py +++ b/pandas/tests/io/test_stata.py @@ -1297,3 +1297,11 @@ def test_pickle_path_localpath(self): reader = lambda x: read_stata(x).set_index('index') result = tm.round_trip_localpath(df.to_stata, reader) tm.assert_frame_equal(df, result) + + def test_read_index(self): + df = tm.makeDataFrame() + df.index.name = "my_index" + with tm.ensure_clean() as path: + df.to_stata(path) + result = pd.read_stata(path, index_col="my_index") + tm.assert_frame_equal(df, result)