如何避免用if-else检查Python函数参数？优化Mongo入库代码

问题场景 #

我写了一个将DataFrame数据保存到MongoDB的Python函数：

def save_to_mongo(self, df, collection, additional_variable):
    for index, row in df.iterrows():
        result = row.to_dict()
        collection.update_one(
            {"_id": str(row['date']) + "_" + str(row['sector']) + "_" + str(row['caseid']) + str(row[additional_variable])},
            {'$set': result},
            upsert=True
        )

但additional_variable参数可能为None，如果用冗长的if-else来处理两种情况，代码会变得丑陋且难以维护：

if additional_variable is None:
    collection.update_one(
        {"_id": str(row['date']) + "_" + str(row['sector']) + "_" + str(row['caseid'])},
        {'$set': result},
        upsert=True
    )
else:
    collection.update_one(
        {"_id": str(row['date']) + "_" + str(row['sector']) + "_" + str(row['caseid']) + str(row[additional_variable])},
        {'$set': result},
        upsert=True
    )

请问有没有更好的方法避免这类重复代码？

解决方案 #

这种重复的分支逻辑确实会让代码冗余，下面是几个更优雅的处理方式：

1. 动态构建_id组件列表

通过先收集必要的_id部分，再根据可选参数动态添加内容，最后统一拼接，避免重复编写update_one逻辑：

def save_to_mongo(self, df, collection, additional_variable=None):
    for index, row in df.iterrows():
        result = row.to_dict()
        # 初始化必填的_id组件
        id_parts = [
            str(row['date']),
            str(row['sector']),
            str(row['caseid'])
        ]
        # 若可选参数存在，追加对应的值到组件列表
        if additional_variable is not None:
            id_parts.append(str(row[additional_variable]))
        # 用下划线拼接所有组件生成最终_id
        doc_id = "_".join(id_parts)
        
        collection.update_one(
            {"_id": doc_id},
            {'$set': result},
            upsert=True
        )

这种方式可读性强，逻辑清晰，后续要修改必填字段或添加新的可选字段时，只需调整id_parts列表即可。

2. 使用f-string结合条件表达式

如果追求代码简洁，可以用f-string的内嵌条件表达式来直接拼接_id，一行搞定：

def save_to_mongo(self, df, collection, additional_variable=None):
    for index, row in df.iterrows():
        result = row.to_dict()
        # 通过条件表达式动态拼接可选部分
        doc_id = f"{row['date']}_{row['sector']}_{row['caseid']}{f'_{row[additional_variable]}' if additional_variable is not None else ''}"
        
        collection.update_one(
            {"_id": str(doc_id)},
            {'$set': result},
            upsert=True
        )

注意这里要确保row[additional_variable]的值可以被字符串化，避免报错。

3. 提取_id生成逻辑为独立辅助函数

如果有多个类似的函数都需要生成类似结构的_id，最推荐的方式是把_id生成逻辑抽成单独的工具函数，彻底消除重复代码：

def generate_doc_id(row, required_fields, additional_field=None):
    """生成MongoDB文档的_id"""
    id_parts = [str(row[field]) for field in required_fields]
    if additional_field is not None:
        id_parts.append(str(row[additional_field]))
    return "_".join(id_parts)

def save_to_mongo(self, df, collection, additional_variable=None):
    for index, row in df.iterrows():
        result = row.to_dict()
        doc_id = generate_doc_id(
            row,
            required_fields=['date', 'sector', 'caseid'],
            additional_field=additional_variable
        )
        
        collection.update_one(
            {"_id": doc_id},
            {'$set': result},
            upsert=True
        )

这样后续任何需要生成同类_id的函数都可以复用generate_doc_id，修改规则时只需改动这一个函数，维护成本大幅降低。

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