Writing function including pandas query with numeric value in function cal-CodePudding

I'm trying to write a function with two calls, one which is the data frame and the other which is some numeric value. I get the error that name "t is not defined." When I hard code that numeric value, everything works well.

Here is minimal reproducible example.

df = pd.DataFrame([[1, 2], [1, 3], [4, 6]], columns=['A', 'B'])

def my_fun(t, df):
    l = df.query('A == t')
    return l

my_fun(1, df)

CodePudding user response：

You could use an f-string to replace the value of the variable t in your query string instead of a literal "t" string:

l = df.query(f"A == {t}")

Complete code:

df = pd.DataFrame([[1, 2], [1, 3], [4, 6]], columns=['A', 'B'])

def my_fun(t, df):
    l = df.query(f'A == {t}')
    return l

print(my_fun(1, df))

Output:

   A  B
0  1  2
1  1  3

CodePudding user response：

You can also use @:

def my_fun(t, df):
    l = df.query('A == @t')
    return l

From the docs:

You can refer to variables in the environment by prefixing them with an ‘@’ character like @a b.