如何根据代码中变量名自动设置Matplotlib绘图的轴标签？

当然可以实现这个需求！我们可以借助Python的inspect模块，获取变量对应的名称字符串，从而自动为绘图设置轴标签。下面是具体的实现方案：

完整代码示例 #

import matplotlib.pyplot as plt
import numpy as np
import inspect

def set_axis_labels(x, y):
    # 获取调用该函数的上下文环境
    caller_frame = inspect.currentframe().f_back
    # 通过变量值匹配找到对应的变量名
    x_label = next(name for name, val in caller_frame.f_locals.items() if val is x)
    y_label = next(name for name, val in caller_frame.f_locals.items() if val is y)
    # 设置轴标签
    plt.xlabel(x_label)
    plt.ylabel(y_label)

# 你的示例代码
x_variable = np.linspace(1, 5)
y_variable = x_variable**2

plt.plot(x_variable, y_variable)
set_axis_labels(x_variable, y_variable)  # 自动设置轴标签
plt.show()

工作原理 #

• inspect.currentframe().f_back：获取调用set_axis_labels函数的上层帧对象，也就是定义x_variable和y_variable的上下文。
• 通过遍历帧对象的局部变量字典f_locals，匹配变量的值（用is判断身份，确保是同一个对象），找到对应的变量名称字符串。
• 最后调用plt.xlabel和plt.ylabel自动设置标签，效果完全等价于手动写plt.xlabel('x_variable')和plt.ylabel('y_variable')。

注意事项 #

• 如果存在多个变量指向同一个数组对象（比如x_alias = x_variable），这种情况下会返回第一个匹配到的变量名，使用时要注意变量的唯一性。
• 这个方法仅适用于变量存储的数组，如果你直接传入表达式（比如plt.plot(x_variable, x_variable**2)），因为表达式没有对应的变量名，无法自动获取标签。

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