Merge branch 'dev' of https://github.com/usc-isi-i2/rltk into dev

usc-isi-i2 · Nov 28, 2018 · ddb8ab1 · ddb8ab1
2 parents e3469ea + b37430d
commit ddb8ab1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/rltk/io/reader/dataframe_reader.py b/rltk/io/reader/dataframe_reader.py
@@ -15,5 +15,5 @@ def __init__(self, df: pd.DataFrame):
         self._df = df
 
     def __next__(self):
-        for _, item in self._df.iterrows():
-            yield item.to_dict()
+        for i, item in self._df.iterrows():
+            yield dict(item.to_dict(), dataframe_default_index=i)
diff --git a/rltk/tests/test_io_reader.py b/rltk/tests/test_io_reader.py
@@ -17,7 +17,7 @@ def test_array_reader():
 def test_dataframe_reader():
     df = pd.DataFrame(arr)
     for idx, obj in enumerate(DataFrameReader(df)):
-        assert obj == arr[idx]
+        assert obj == dict(**arr[idx], dataframe_default_index=idx)
 
 
 def test_csv_reader():
@@ -43,4 +43,4 @@ def test_jsonlines_reader():
     for idx, obj in enumerate(JsonLinesReader(f)):
         assert obj == arr[idx]
 
-    f.close()
+    f.close()