BUG: Preserve columns in cum_returns.

- Fixes a bug where we failed to maintain columns when calling `empyrical.cum_returns` on a DataFrame.
quantopian · Mar 12, 2018 · 708ee2a · 708ee2a
1 parent 1c0978d
commit 708ee2a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/empyrical/stats.py b/empyrical/stats.py
@@ -237,7 +237,9 @@ def cum_returns(returns, starting_value=0, out=None):
         if returns.ndim == 1 and isinstance(returns, pd.Series):
             out = pd.Series(out, index=returns.index)
         elif isinstance(returns, pd.DataFrame):
-            out = pd.DataFrame(out, index=returns.index)
+            out = pd.DataFrame(
+                out, index=returns.index, columns=returns.columns,
+            )
 
     return out
 

diff --git a/empyrical/tests/test_stats.py b/empyrical/tests/test_stats.py
@@ -36,6 +36,10 @@ def assert_indexes_match(self, result, expected):
         """
         assert_index_equal(result.index, expected.index)
 
+        if isinstance(result, pd.DataFrame) and \
+           isinstance(expected, pd.DataFrame):
+            assert_index_equal(result.columns, expected.columns)
+
 
 class TestStats(BaseTestCase):