def manipulate(self,
res: pd.Series,
melted: bool = False,
return_dataframe: bool = True):
"""Common adhoc data manipulation.
It does
1. Converts res to a DataFrame if asked.
2. Melts res to long format if asked.
3. Removes redundant index levels in res.
Args:
res: Returned by compute_through(). Usually a DataFrame, but could be a
pd.Series or a base type.
melted: Whether to transform the result to long format.
return_dataframe: Whether to convert the result to DataFrame if it's not.
If False, it could still return a DataFrame if the input is already a
DataFrame.
Returns:
Final result returned to user. If split_by, it's a pd.Series or a
pd.DataFrame, otherwise it could be a base type.
"""
if isinstance(res, pd.Series):
res.name = self.name
res = self.to_dataframe(res) if return_dataframe else res
if melted:
res = utils.melt(res)
return utils.remove_empty_level(res)
def test_remove_empty_level(self):
df = pd.DataFrame([[0, 1, 2]], columns=['a', 'b', 'c'])
df.set_index(['b', 'a'], append=True, inplace=True)
expected = df.droplevel(0)
actual = utils.remove_empty_level(df)
testing.assert_frame_equal(expected, actual)