fix: read --columns filters output.

timcera · Nov 23, 2020 · 7b8d137 · 7b8d137
1 parent a3edf4c
commit 7b8d137
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/tstoolbox/functions/read.py b/tstoolbox/functions/read.py
@@ -147,7 +147,6 @@ def read(
             ),
             start_date=start_date,
             end_date=end_date,
-            pick=columns,
             round_index=round_index,
             dropna=dropna,
             force_freq=force_freq,
@@ -171,6 +170,8 @@ def read(
     if append != "combine":
         result = pd.concat(result_list, axis=append)
 
+    result = tsutils._pick(result, columns)
+
     result.sort_index(inplace=True)
 
     return result

diff --git a/tstoolbox/tsutils.py b/tstoolbox/tsutils.py
@@ -1292,7 +1292,10 @@ def _pick(tsd, columns):
             # Use the -1 marker to indicate index
             jtsd = pd.DataFrame(tsd.index, index=tsd.index)
         else:
-            jtsd = pd.DataFrame(tsd[tsd.columns[col]], index=tsd.index)
+            try:
+                jtsd = pd.DataFrame(tsd.iloc[:, col], index=tsd.index)
+            except IndexError:
+                jtsd = pd.DataFrame(tsd.loc[:, col], index=tsd.index)
 
         newtsd = newtsd.join(jtsd, lsuffix="_{0}".format(index), how="outer")
     return newtsd