Python函数基于条件返回不同变量的可行性及解包报错解决

首先，你遇到的ValueError: too many values to unpack (expected 2)报错，核心原因是函数返回值和你试图接收的变量数量不匹配：

当你设置subset=True时，你期望函数返回两个子集，但你的函数最后始终只返回了embedded_df这一个DataFrame，自然没法用两个变量去接收它。另外，你的代码里还有一个隐藏bug：内部函数subsetting里用到了未定义的变量gal_subset，这会导致调用subset=True时直接触发NameError，只是当前的返回值问题先暴露出来了。

修正后的函数代码 #

import pandas as pd

def modify_data(embeddings, df, subset=False):
    """ Modifies Original Dataframe with respective embeddings 
    :return: 
        - If subset=False: Final processed Dataframe
        - If subset=True: Tuple of two subset Dataframes (first 1000 rows of 2000-sampled data, remaining 1000 rows)
    """
    # Original_DF processing
    OD_df = df.copy(deep=True)
    OD_df = OD_df.reset_index()
    OD_df.loc[:, 'task'] = 'stability'
    
    # Embeddings Dataframe processing
    embeddings_df = pd.DataFrame(data=embeddings)
    embeddings_df = embeddings_df.reset_index()
    embedded_df = pd.merge(embeddings_df, OD_df, on='index')
    embedded_df = embedded_df.drop(['index', 'sequence', 'temperature'], axis=1)
    
    def subsetting(df_to_subset, sample_no, row_no):
        "Select a subset of rows from the processed dataframe"
        # First sample the desired number of rows
        sampled_df = df_to_subset.sample(n=sample_no)
        # Split into two subsets
        subset_first = sampled_df[:row_no]
        subset_second = sampled_df[row_no:]
        return subset_first, subset_second
    
    # 根据subset参数返回对应结果
    if subset:
        # 返回两个子集组成的元组
        return subsetting(embedded_df, sample_no=2000, row_no=1000)
    else:
        # 返回完整的处理后DataFrame
        return embedded_df

关键修改点说明 #

1. 修复未定义变量问题：把subsetting里的gal_subset改成了sampled_df（先把采样后的结果存到这个变量，再进行切片），避免NameError。
2. 调整返回逻辑：根据subset参数的取值，返回不同类型的结果：
   • 当subset=False时，返回处理后的完整DataFrame，和你原来的逻辑一致。
   • 当subset=True时，直接返回subsetting函数生成的两个子集组成的元组，这样你就能用两个变量去接收它。
3. 更新文档字符串：明确标注不同参数下的返回值类型，让函数的行为更清晰，后续维护也更方便。

正确调用方式 #

• 当不需要子集时：

  modified_df = modify_data(protein_embeddings, protein_df, subset=False)
  

• 当需要子集时：

  gal_subset_first, gal_subset_second = modify_data(protein_embeddings, protein_df, subset=True)
  

这样修改后，两种调用场景都能正常工作啦～

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