yaml.load，强制将字典键转为字符串

Question

在Python3中，我正在加载一个YAML片段。 加载器试图猜测正确的类型，但我不太满意。 我希望字典键始终是字符串。

首先是一个最小化的YAML片段，可直接粘贴到您的Python解释器中进行演示。 不用说，我的现实世界数据要复杂得多。

txt = """
---
one: 1
2: two
"""

首先是“常规”负载：

yaml.load(txt)
{2: 'two', 'one': 1}

注意键2被加载为数字而不是字符串。现在让我们尝试一些不同的东西：

yaml.load(txt, Loader=yaml.BaseLoader)
{'2': 'two', 'one': '1'}

现在所有的东西都被作为字符串处理。不幸的是，1 也是字符串，但我需要它作为一个数字值。所以我要么把键和值都强制转化为字符串，要么不转化。

当然，我可以创建一个后处理程序来遍历加载的数据并将其复制到一个新变量中，将字典键强制转化为字符串，但我认为这可以更优雅地在YAML加载器内部完成。有什么建议吗？

- mogul

Answer 1

您可以用几行代码完成此操作，将正在构建的每个映射更改为在运行时将整数类型键转换为字符串。您可以子类化SafeLoader，但然后需要注册构造函数。最简单的方法是修补映射构造函数：

import yaml

def my_construct_mapping(self, node, deep=False):
    data = self.construct_mapping_org(node, deep)
    return {(str(key) if isinstance(key, int) else key): data[key] for key in data}

yaml.SafeLoader.construct_mapping_org = yaml.SafeLoader.construct_mapping
yaml.SafeLoader.construct_mapping = my_construct_mapping


yaml_str = """\
---
one: 1
2: two
"""

data = yaml.safe_load(yaml_str)
print(data)

这将给出：

{'one': 1, '2': 'two'}

---

永远不要使用默认的、不安全的 yaml.load()（即没有 Loader= 参数）。

- Anthon

---

谢谢，这个肯定有效。在我看来有点像猴子补丁，但是好吧，它完成了工作。至少它给了我进一步开发的东西。我担心我将不得不做一个适当的子类化 - 在某些时候我可能需要一个“本地”的yaml加载器，以及你的“黑客”加载器。 - mogul

@mogul “本地”加载器涉及对SafeConstructor进行子类化，在该子类上多次调用.add_constructor()（参见constructor.py），然后创建一个基于SafeLoader的加载器，该加载器使用该子类进行构造。使用我的ruamel.yaml库比PyYAML更容易实现，但仍需要大量代码。如果您遇到困难，请知道在哪里发布问题;-) - Anthon

Answer 2

被接受的答案对浮点数无效（我不知道这是否是您的意图），并且它也会将布尔值更改为字符串。而且这是一个猴子补丁。

要以正确的方式执行此操作，您需要对加载器中的一个进行子类化，例如yaml.SafeLoader：

import yaml


class MyLoader(yaml.SafeLoader):
    def construct_mapping(self, *args, **kwargs):
        mapping = super().construct_mapping(*args, **kwargs)

        for key in list(mapping.keys()):
            # bool is a subclass of int
            if not isinstance(key, bool) and isinstance(key, (int, float)):
                mapping[str(key)] = mapping.pop(key)

        return mapping

txt = """
---
one: 1
2: two
3.0: three
"""

yaml.load(txt, Loader=MyLoader)

这将给你：

{'one': 1, '2': 'two', '3.0': 'three'}

如果您不希望它适用于浮点数，请将(int, float)替换为int。

请注意，如果一个键是整数，另一个键是字符串，它将删除键：

txt = """
1: one
'1': another one
"""
yaml.load(txt, Loader=MyLoader)

会给：

{'1': 'one'}

- FunkyPotato

---

如果您不想进行猴子补丁，那么降级到PyYAML是没有必要的。它仅支持YAML 1.1规范的子集，而该规范已于2009年被取代。因此，在我看来，这样做只会制造更多问题而不是解决问题。 - Anthon