如何扩展PyYaml/ruamel.yaml Loader，将未引号inf/nan解析为float？

一、PyYAML 实现方式 #

要实现未加引号的inf/nan解析为float类型、加引号的保留字符串的逻辑，核心是扩展PyYAML的Loader，修改标量解析规则：

import yaml
from yaml.resolver import Resolver

class InfNanLoader(yaml.SafeLoader):
    pass

# 匹配不区分大小写的未加引号inf、nan（含正负号）
inf_nan_regex = r'|([+-]?inf|INF)|([+-]?nan|NAN)$'
# 扩展默认float解析正则，将inf/nan纳入匹配范围
original_float_re = yaml.resolver.FloatResolver.DEFAULT_REGEX.pattern
new_float_re = original_float_re.replace('|$', inf_nan_regex)

Resolver.add_implicit_resolver(
    u'tag:yaml.org,2002:float',
    new_float_re,
    list('-+0123456789.infINFnNaN')
)

# 测试验证
test_cases = [
    '[inf, nan]',
    "['inf', 'nan']",
    '[Inf, NaN]',
    '[+inf, -nan]'
]

for case in test_cases:
    result = yaml.load(case, Loader=InfNanLoader)
    print(f"输入: {case}")
    print(f"输出: {result}")
    print(f"元素类型: {[type(x) for x in result]}")
    print('---')

这段代码通过扩展float类型的隐式解析规则，让未加引号的inf/nan触发float解析，加引号的字符串则不受影响。

二、ruamel.yaml 实现方式 #

ruamel.yaml的扩展逻辑类似，需修改默认的标量正则映射：

import re
from ruamel.yaml import YAML
from ruamel.yaml.resolver import Resolver

yaml = YAML(typ='safe')
# 扩展float解析正则，加入inf/nan匹配
original_float_re = Resolver.DEFAULT_SCALAR_TAG_TO_REGEX[u'tag:yaml.org,2002:float']
new_float_re = original_float_re.pattern.replace('|$', r'|([+-]?inf|INF)|([+-]?nan|NAN)$')
Resolver.DEFAULT_SCALAR_TAG_TO_REGEX[u'tag:yaml.org,2002:float'] = re.compile(new_float_re)

# 测试验证
test_cases = [
    '[inf, nan]',
    "['inf', 'nan']",
    '[Inf, NaN]',
    '[+inf, -nan]'
]

for case in test_cases:
    result = yaml.load(case)
    print(f"输入: {case}")
    print(f"输出: {result}")
    print(f"元素类型: {[type(x) for x in result]}")
    print('---')

三、此类自定义解析的隐患 #

即使仅用于命令行参数解析，仍需注意以下风险：

• 兼容性问题：标准YAML未将未加引号的inf/nan列为合法float标量，其他YAML解析器可能无法识别该写法，导致跨工具解析失败。
• 歧义风险：若存在名为inf的合法字符串参数（如标识符），未加引号会被误解析为float，引发逻辑错误。
• 安全风险：自定义Loader若扩展不当，可能引入注入漏洞（基于SafeLoader扩展可降低该风险）。

四、替代方案（可选） #

如果担心隐患，可优先选择标准兼容写法：

• 使用.inf/.nan：[.inf, .nan]
• 使用显式类型标记：[!!float inf, !!float nan]

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