How can I use a custom function within an expression using the eval dataframe method?

Viewed 1639

I am using Python 3.X.

With the builtin function eval() you can use a dictionaty of objects in order to use a custom function like this:

from math import *

def one():
    # some operations
    return 1

functions = {
    '__builtins__': None,
    'sqrt': sqrt,
    'one': one,
}
variables = {
    '__builtins__': None,
    'pi': pi,
}
expression = 'sqrt(34 * pi) + one()'
eval(expression, variables, functions)

But the eval() dataframe method does not work like that. You can only use these built-in functions:

The supported math functions are sin, cos, exp, log, expm1, log1p, sqrt, sinh, cosh, tanh, arcsin, arccos, arctan, arccosh, arcsinh, arctanh, abs and arctan2

import pandas as pd
import numpy as np
from math import *

df = pd.DataFrame({
    'A': [0, 10, 0, 10, 10, 30],
    'B': [0, 0, 1000, 1000, 0, 0],
    'C': [25, 25, 25, 25, 40, 40]
})

def custom():
    # some operations
    return 3

functions = {
    'custom': custom
}
variables = {
    'pi': pi
}
equation = 'D = sqrt(A) + B + custom()'
df.eval(
    equation, global_dict=variables, local_dict=functions,
    engine='numexpr', inplace=True
)
# ERROR: "custom" is not a supported function

Is there a way to use a custom function in the expression?

NOTE: I know it could bedangerous, but it is on me

1 Answers

Use @ when calling local variables or local functions:

In [45]: equation = 'D = sqrt(A) + B + @custom()'
#  NOTE:   ------------>               ^

In [46]: df.eval(equation, inplace=True)

In [47]: df
Out[47]:
    A     B   C            D
0   0     0  25     3.000000
1  10     0  25     6.162278
2   0  1000  25  1003.000000
3  10  1000  25  1006.162278
4  10     0  40     6.162278
5  30     0  40     8.477226
Related