如何将自定义类传入Matplotlib进行绘图？

方案1：列表推导式直接提取字段 #

无需修改自定义类，调用plt.plot时通过列表推导式直接提取time字段，代码简洁直观：

from dataclasses import dataclass
import matplotlib.pyplot as plt

@dataclass
class TestCase:
    name: str
    time: float

# 实例化测试用例
d1 = TestCase("Case1", 2.5)
d2 = TestCase("Case2", 4.3)
d3 = TestCase("Case3", 1.8)

# 提取time字段绘图
plt.plot([tc.time for tc in [d1, d2, d3]])
plt.ylabel('Execution Time')
plt.show()

方案2：给自定义类添加数值转换方法 #

如果想实现plt.plot([d1, d2])这种直接传入实例列表的调用方式，给TestCase类添加__float__方法，让Matplotlib自动将实例转换为对应的time数值：

from dataclasses import dataclass
import matplotlib.pyplot as plt

@dataclass
class TestCase:
    name: str
    time: float

    def __float__(self):
        return self.time

# 实例化测试用例
d1 = TestCase("Case1", 2.5)
d2 = TestCase("Case2", 4.3)
d3 = TestCase("Case3", 1.8)

# 直接传入实例列表绘图
plt.plot([d1, d2, d3])
plt.ylabel('Execution Time')
plt.show()

方案3：使用map函数提取字段 #

习惯函数式编程的话，可以用map函数结合lambda表达式提取time字段：

# 基于方案1的类定义和实例
plt.plot(list(map(lambda tc: tc.time, [d1, d2, d3])))
plt.ylabel('Execution Time')
plt.show()

内容的提问来源于stack exchange，提问作者Matthieu Raynaud de Fitte