How to include pandas comparison operators in the function?-CodePudding

Let's say I have the following data

import pandas as pd 

d = {'index': [0, 1, 2], 'a': [10, 8, 6], 'b': [4, 2, 6],}

data_frame  = pd.DataFrame(data=d).set_index('index')

Now what I do is, filter this data based on the values of column "b", Lets say like this:

new_df = data_frame[data_frame['b']!=4]
new_df1 = data_frame[data_frame['b']==4]

What I want to do, instead of this method above, is to write the function where I can also indicate what kind of comparison operators it should use. Something like this


def slice(df, column_name):
    df_new = df[df[column_name]!=4]
    return df_new

new_df2 = slice(df=data_frame, column_name='b')

The function above only does != operation in the data. What I want is that both != and == to be somehow defined in the above function, and when I use the function I could ideally indicate which one to use.

Let me know if my question needs more detailed clarification

CodePudding user response：

You could add a boolean parameter to your function:

def slice(df, column_name, equality=True):
    if equality:
        df_new = df[df[column_name]==4]
    else:
        df_new = df[df[column_name]!=4]
    return df_new

new_df2 = slice(df=data_frame, column_name='b', equality=True)

By the way, slice is a built-in python function, so probably a good idea to rename to something else.