Python中str.match正则在主函数可用，移入其他函数报错

问题背景 #

需要处理格式为AANNAA（A代表字母、N代表数字）的6位编码，给符合模式的行标记0，不符合的标记1。主函数内的代码能将不符合的编码识别为NaN，但移至自定义函数后抛出ValueError: Cannot set a DataFrame with multiple columns to the single column code错误。

主函数中可运行的代码片段：

wave2['code'] = wave2[wave2['SGIC2'].str.match(r'^[A-Za-z]{2}[0-9]{2}[A-Za-z]{2}$') ==True]

完整报错代码：

import pandas as pd

wave2 = pd.read_csv ('wave2.csv')

def main():
    wave2.rename(columns={ wave2.columns[0]: "SGIC2" }, inplace = True)

    #wave2['code'] = wave2[wave2['SGIC2'].str.match(r'^[A-Za-z]{2}[0-9]{2}[A-Za-z]{2}$') ==True]
    #print(wave2)
    check_valid(wave2)

def check_valid(dfName):          #check length and composition AANNAA
    length = (dfName.iloc[:,0].str.len())
    dfName['lengthinV'] = (length != 6).astype(int)
    nlengthinV = str(dfName['lengthinV'].values.sum())    #will not be needed once regex works

    dfName['code'] = dfName[dfName.iloc[:,0].str.match(r'^[A-Za-z]{2}[0-9]{2}[A-Za-z]{2}$')==True]
    print(dfName)
    return nlengthinV

if __name__ == "__main__":
    main()

Data = ('AB01ER','DA23RE','MN34ER','FG19SD','BB21BB', 'TR15HG','SE21AR','TI85BV','LK31YU','WI29VV','WI13AL', 'HL29WE','HL29WE','IH8THS','TH15P8', 'AS43GGG', 'J12RT',  'RT13CA', 'CH08VI', 'KK09DE')

期望输出：

SGIC2    code  lengthinV
0    AB01ER  AB01ER          0
1    DA23RE  DA23RE          0
2    MN34ER  MN34ER          0
3    FG19SD  FG19SD          0
4    BB21BB  BB21BB          0
5    TR15HG  TR15HG          0
6    SE21AR  SE21AR          0
7    TI85BV  TI85BV          0
8    LK31YU  LK31YU          0
9    WI29VV  WI29VV          0
10   WI13AL  WI13AL          0
11   HL29WE  HL29WE          0
12   HL29WE  HL29WE          0
13   IH8THS     NaN          0
14   TH15P8     NaN          0
15  AS43GGG     NaN          1
16    J12RT     NaN          1
17   RT13CA  RT13CA          0
18   CH08VI  CH08VI          0
19   KK09DE  KK09DE          0

错误原因 #

主函数中执行赋值时，DataFrame仅包含SGIC2一列，因此wave2[条件]返回的是单列DataFrame，赋值给code列不会报错。但在自定义函数中，已经添加了lengthinV列，此时dfName[条件]返回的是包含SGIC2和lengthinV的多列DataFrame，尝试将多列数据赋值给单个code列，就会触发上述ValueError。

解决方案 #

修改自定义函数中的赋值逻辑，仅提取符合正则条件的SGIC2列值，不符合的自动设为NaN。推荐使用where方法实现：

修正后的check_valid函数：

def check_valid(dfName):          #check length and composition AANNAA
    # 提取第一列（即SGIC2列）
    sgic_col = dfName.iloc[:,0]
    length = sgic_col.str.len()
    dfName['lengthinV'] = (length != 6).astype(int)
    nlengthinV = str(dfName['lengthinV'].values.sum())

    # 仅将符合正则的SGIC2值赋值给code列，不符合的设为NaN
    dfName['code'] = sgic_col.where(sgic_col.str.match(r'^[A-Za-z]{2}[0-9]{2}[A-Za-z]{2}$'))
    print(dfName)
    return nlengthinV

也可以使用布尔索引提取对应列的方式：

dfName['code'] = dfName.loc[sgic_col.str.match(r'^[A-Za-z]{2}[0-9]{2}[A-Za-z]{2}$'), dfName.columns[0]]

这种写法会自动将不符合条件的行设为NaN，同样能达到需求。

验证结果 #

运行修正后的代码，即可得到与期望一致的输出：符合AANNAA格式的编码会复制到code列，不符合的显示NaN，lengthinV列正确标记编码长度是否为6。

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