mardi 30 mars 2021

Python - reduce number of if statements

I am filtering a pandas dataframe based on one or more conditions, like so:

def filter_dataframe(dataframe, position=None, team_id=None, home=None, window=None, min_games=0):
        
        df = dataframe.copy()

        if position:
            df = df[df['posicao_id'] == position] 
        
        if clube_id:
            df = df[df['team_id'] == team_id]
        
        if home:
            if home == 'home':
                df = df[df['home_dummy'] == 1.0]
            elif home == 'away':
                df = df[df['home_dummy'] == 0.0]
        
        if window:
            df = df[df['round_id'].between(1, window)]
        
        if min_games:
            df = df[df['games_num'] >= min_games]

        return df

But I don't think this is elegant.

Is there a simpler way of achieving the same result?

I though of creating rules for conditions like in this SO answer and then use the method any(rules) in order to apply the filtering, if any, but I don't know how to approach this. Any ideas?

Aucun commentaire:

Enregistrer un commentaire