如何在Python中将数组变量名转为字符串用作图表标题？

问题背景 #

想要让函数自动把输入变量名作为图表标题，尝试用列表推导式从locals()中匹配变量值来获取变量名，但运行时报错。这种方法对单值变量（如整数）有效，但对数组类型的输入会触发ValueError，原因是数组元素过多导致比较结果的真值判断歧义。

原代码 #

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from scipy.fftpack import hilbert, ihilbert

def mod_arg_plot(column, time):
    # 尝试获取变量名的代码，注释后可运行
    # title = [i for i, j in locals().items() if j == column][0]

    plt.plot(time, abs(ihilbert(column)), 'b-', time, np.angle(ihilbert(column)), 'r--')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude / degrees')
    plt.legend(('Modulus', 'Argument'))
    # plt.title(title)
    plt.show()

报错信息 #

line 44, in mod_arg_plot
    title = [i for i, j in locals().items() if j == column][0]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
line 44, in <listcomp>
    title = [i for i, j in locals().items() if j == column][0]
                                               ^^^^^^^^^^^
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

问题原因 #

当column是NumPy数组时，j == column会返回一个布尔数组（逐元素比较的结果），而if条件需要单个布尔值来判断，因此Python抛出歧义错误。另外，即使两个数组内容完全相同，它们也可能是不同的对象，用==匹配值的方式本身就不可靠。

解决方案 #

方案1：显式传递标题参数（最稳定可靠） #

直接在调用函数时传入标题，避免依赖变量名的自动推导，这种方式最清晰，也能处理任何输入场景（比如传递的是计算后的数组、切片等没有明确变量名的情况）。

修改后的代码：

import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import ihilbert

def mod_arg_plot(column, time, title=None):
    # 如果没传标题，默认用'Column Data'
    plt_title = title if title is not None else 'Column Data'
    
    plt.plot(time, abs(ihilbert(column)), 'b-', time, np.angle(ihilbert(column)), 'r--')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude / degrees')
    plt.legend(('Modulus', 'Argument'))
    plt.title(plt_title)
    plt.show()

# 调用示例
my_data = np.random.randn(100)
my_time = np.linspace(0, 10, 100)
mod_arg_plot(my_data, my_time, title='my_data')

方案2：利用inspect模块自动获取调用时的变量名 #

如果确实需要自动获取变量名，可以用inspect模块解析调用栈，提取传入参数的表达式。这种方式能自动获取简单变量名，但如果传入的是复杂表达式（比如my_data[::2]），得到的会是表达式字符串而非变量名。

修改后的代码：

import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import ihilbert
import inspect
import ast

def mod_arg_plot(column, time):
    # 获取调用栈，解析传入column的参数表达式
    frame = inspect.currentframe().f_back
    call_node = ast.walk(ast.parse(inspect.getframeinfo(frame).code_context[0])).__next__()
    # 找到column对应的参数值
    for arg in call_node.args:
        if isinstance(arg, ast.Name):
            title = arg.id
            break
    else:
        title = 'Column Data'

    plt.plot(time, abs(ihilbert(column)), 'b-', time, np.angle(ihilbert(column)), 'r--')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude / degrees')
    plt.legend(('Modulus', 'Argument'))
    plt.title(title)
    plt.show()

# 调用示例
my_data = np.random.randn(100)
my_time = np.linspace(0, 10, 100)
mod_arg_plot(my_data, my_time)  # 标题会自动设为'my_data'

注意事项 #

• 方案2的自动推导仅适用于简单变量传递场景，若传入的是表达式、匿名函数结果等，无法得到预期的变量名。
• 依赖变量名的方式本质上不可靠，因为变量名只是引用，同一个值可能被多个变量引用，或者传递的是没有变量名的临时对象。

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