Python迭代器函数未遍历列表全部元素的技术求助

问题描述 #

编写了两个Python函数用于提取化学式中的原子数量，但iterator函数仅处理列表第一个元素，无法正确找到目标原子（如"O"）的数量：

原函数代码 #

import re

# function to get number of wanted atom
def atom_number_grabber(sum_formula, wanted_atom):
    match = re.match(r"([A-Z][a-z]*)([0-9]*)", sum_formula, re.I)
    if match:
        items = match.groups()
    if items[0] == wanted_atom:
        atom_number = items[1]
        if not atom_number:
            atom_number = "1"
        return atom_number
    else:
        pass

# function to iterate over all elements
def iterator(sum_formula_list, atom_number_grabber, wanted_atom):
    for sum_formula in sum_formula_list:
        return atom_number_grabber(sum_formula, wanted_atom)

测试代码与结果 #

test_list = ["C25", "H32", "O8"]

iterator(sum_formula_list = test_list, atom_number_grabber = atom_number_grabber, wanted_atom = "O")

• 实际输出：空（返回None）
• 期望输出：8

当目标原子改为"C"时，函数返回25，说明仅处理了列表第一个元素。

问题原因 #

1. iterator函数提前终止循环：for循环的第一次迭代就执行return语句，直接终止循环并返回第一个元素的处理结果，后续元素完全未被遍历。
2. atom_number_grabber存在逻辑漏洞：若正则匹配失败，items变量未定义会抛出UnboundLocalError；未匹配到目标原子时无明确返回值，默认返回None。

修复方案 #

优化atom_number_grabber函数 #

添加边界处理，确保所有分支都有明确返回值：

import re

# function to get number of wanted atom
def atom_number_grabber(sum_formula, wanted_atom):
    match = re.match(r"([A-Z][a-z]*)([0-9]*)", sum_formula, re.I)
    if match:
        items = match.groups()
        if items[0] == wanted_atom:
            atom_number = items[1]
            return atom_number if atom_number else "1"
    # 未匹配到目标原子或格式错误时返回None
    return None

修复iterator函数 #

遍历所有元素，找到目标原子后再返回结果，遍历结束未找到则返回None：

# function to iterate over all elements
def iterator(sum_formula_list, atom_number_grabber, wanted_atom):
    for sum_formula in sum_formula_list:
        result = atom_number_grabber(sum_formula, wanted_atom)
        if result is not None:
            return result
    # 遍历完未找到目标原子时返回None（可根据需求改为0或其他值）
    return None

扩展：收集所有匹配结果 #

如果需要提取列表中所有目标原子的数量（如存在多个相同原子项），可修改iterator返回结果列表：

def iterator(sum_formula_list, atom_number_grabber, wanted_atom):
    results = []
    for sum_formula in sum_formula_list:
        result = atom_number_grabber(sum_formula, wanted_atom)
        if result is not None:
            results.append(result)
    return results

此时原测试用例将返回['8']，可根据需求进一步转换为数值类型。

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