Python函数中如何获取传入参数的变量名并赋值给DataFrame列

问题原因 #

你原代码里的f'{argument1=}'.split('=')[0]只能获取函数内部的参数名argument1，而不是外部传入的实际变量名，所以无法得到期望的string_i_want。

方法1：通过inspect模块解析调用上下文 #

这种方法可以自动提取调用时传入的变量名，但仅适用于直接传递变量的简单场景：

import pandas as pd
import inspect

string_i_want = pd.DataFrame([1, 2, 3, 4], columns=['column_name'])

def add_scenario_col(df):
    # 获取调用栈的上下文帧
    caller_frame = inspect.currentframe().f_back
    # 提取调用函数时的代码文本
    caller_code = inspect.getframeinfo(caller_frame).code_context[0].strip()
    # 解析出传入的变量名（仅支持直接传变量的情况）
    var_name = caller_code.split('(')[1].rstrip(')')
    
    df['squared'] = df['column_name'] ** 2
    df['scenario'] = var_name
    return df

add_scenario_col(string_i_want)

注意：如果调用时传入的是表达式（比如add_scenario_col(pd.read_csv("data.csv"))），这种方法会失效，因为没有明确的变量名可以提取。

方法2：显式传入数据集名称（推荐，更稳定） #

这种方法虽然需要手动传入名称，但完全不依赖调用上下文，兼容性和可靠性更强：

import pandas as pd

string_i_want = pd.DataFrame([1, 2, 3, 4], columns=['column_name'])

def add_scenario_col(df, scenario_name):
    df['squared'] = df['column_name'] ** 2
    df['scenario'] = scenario_name
    return df

add_scenario_col(string_i_want, 'string_i_want')

内容的提问来源于stack exchange，提问作者windyvation