如何构建可根据用户输入生成任意数量国家折线图的Python类

你可以通过可变位置参数*countries接收任意数量的国家入参，遍历所有国家批量绘制折线即可，一套逻辑兼容1到任意数量国家的绘图需求。

优化后完整代码 #

import pandas as pd
import matplotlib.pyplot as plt

class ExploreData:
    def __init__(self, data_path='mydata.csv'):
        # 初始化时读取数据，支持自定义文件路径
        self.df = pd.read_csv(data_path)
        # 提前把日期设为索引，避免每次绘图重复处理
        self.df = self.df.set_index('date')

    def plot_country_trend(self, *countries, var_col='var'):
        """
        绘制任意数量国家的指标趋势折线图
        :param countries: 要绘制的国家名，支持传入1到N个
        :param var_col: 要绘制的指标列名，默认使用var
        """
        # 遍历每个国家绘制折线
        for country in countries:
            country_data = self.df[self.df['country'] == country]
            plt.plot(country_data[var_col], label=country)
        
        # 自动生成标题
        if len(countries) == 1:
            plt.title(f"{countries[0]} {var_col}趋势")
        else:
            plt.title(f"{'、'.join(countries)} {var_col}对比")
        
        plt.legend()
        plt.show()

使用示例 #

# 实例化类
plotter = ExploreData()

# 绘制单个国家
plotter.plot_country_trend("中国")

# 绘制两个国家对比
plotter.plot_country_trend("中国", "美国")

# 绘制任意多个国家
plotter.plot_country_trend("中国", "美国", "日本", "韩国")

核心优化点 #

• 用*countries接收任意数量的国家入参，无需单独维护单国家、双国家的两套方法
• 提前在初始化方法里完成日期索引设置，避免每次绘图重复计算
• 修复原代码中plot2方法未调用实例属性self.df的bug
• 支持自定义要绘制的指标列，灵活性更高
• 自动适配不同数量国家的标题生成逻辑

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