Search code examples
pythonpandasif-statementmethod-chaining

Conditional method chaining in pandas


Is there a simple general way to make a method conditional to an if-statement when using method chaining with pandas?

Mock example:

df = pd.DataFrame({'A':['one', 'two'], 'B':['one', 'two']})

change_to_numeric = False

df = (df
    .query("A == 'one'")
    .replace('one', 1)         # <-- Execute this row only "if change_to_numeric == True"
)

Thank you!


Solution

  • You can use pipe:

    df = pd.DataFrame({'A':['one', 'two'], 'B':['one', 'two']})
    
    change_to_numeric = False
    
    df = (df
        .query("A == 'one'")
        .pipe(lambda d: d.replace('one', 1) if change_to_numeric else d)
    )
    

    output for change_to_numeric = False:

         A    B
    0  one  one
    

    output for change_to_numeric = True:

       A  B
    0  1  1