I would like to get slices of a pandas.DataFrame using an axis argument to specify either rows or columns. That is, I would like the functionality provided by the function _loc_axis in the MWE below. Is there an easy way to do this?
import pandas as pd
def _loc_axis(df, idx, axis):
"""Call loc along proper axis"""
# Axis as integer, error if axis is invalid
axis = {0: 0, 'index': 0,
1: 1, 'columns': 1}[axis]
return (df.loc[idx, :] if axis==0 else
df.loc[:, idx])
df = pd.DataFrame(
data=[[1,2,3], [4,5,6], [7,8,9]],
index=['a','b','c'],
columns=['D','E','F'],
)
print(df)
print(_loc_axis(df, ['a','b'], axis='index'))
print(_loc_axis(df, ['F','D'], axis='columns'))
Produces output:
D E F
a 1 2 3
b 4 5 6
c 7 8 9
D E F
a 1 2 3
b 4 5 6
F D
a 3 1
b 6 4
c 9 7