使用 Python 从 YAML 中检索批量数据

Question

我有一个如下形式的 yaml 文件：

Solution: 
- number of solutions: 1
  number of solutions displayed: 1
- Gap: None
  Status: optimal
  Message: bonmin\x3a Optimal
  Objective:
    objective:
      Value: 0.010981105395
  Variable:
    battery_E[b1,1,1]:
      Value: 0.25
    battery_E[b1,1,2]:
      Value: 0.259912707017
    battery_E[b1,2,1]:
      Value: 0.120758408109
    battery_E[b2,1,1]:
      Value: 0.0899999972181
    battery_E[b2,2,3]:
      Value: 0.198967393893
    windfarm_L[w1,2,3]:
      Value: 1
    windfarm_L[w1,3,1]:
      Value: 1
    windfarm_L[w1,3,2]:
      Value: 1

使用 Python27，我想从这个 YAML 文件中导入所有的 battery_E 值。我知道我可以遍历 battery_E 字典的键以逐一检索它们（我已经在使用 PyYAML 进行此操作），但我想避免迭代并一次性完成！

Answer 1

不可能“一口气” - 无论哪种方式，都会有某种迭代，这完全没问题。

但是，如果内存是一个问题，您可以在 YAML 加载期间加载仅感兴趣的键的值：

from __future__ import print_function

import yaml

KEY = 'battery_E'


class Loader(yaml.SafeLoader):
    def __init__(self, stream):
        super(Loader, self).__init__(stream)
        self.values = []

    def compose_mapping_node(self, anchor):
        start_event = self.get_event()
        tag = start_event.tag
        if tag is None or tag == '!':
            tag = self.resolve(yaml.MappingNode, None, start_event.implicit)
        node = yaml.MappingNode(tag, [],
                                start_event.start_mark, None,
                                flow_style=start_event.flow_style)
        if anchor is not None:
            self.anchors[anchor] = node
        while not self.check_event(yaml.MappingEndEvent):
            item_key = self.compose_node(node, None)
            item_value = self.compose_node(node, item_key)
            if (isinstance(item_key, yaml.ScalarNode)
                    and item_key.value.startswith(KEY)
                    and item_key.value[len(KEY)] == '['):
                self.values.append(self.construct_object(item_value, deep=True))
            else:
                node.value.append((item_key, item_value))
        end_event = self.get_event()
        node.end_mark = end_event.end_mark
        return node


with open('test.yaml') as f:
    loader = Loader(f)
    try:
        loader.get_single_data()
    finally:
        loader.dispose()

    print(loader.values)

但请注意，此代码不假定 YAML 文件内树中 battery_E 键的位置 - 它只会加载它们的所有值。

Answer 2

无需使用PyYAML检索每个条目，您可以加载一次数据，然后使用Pythons选择以下两行的键值对：

data = yaml.safe_load(open('input.yaml'))
kv = {k:v['Value']  for k, v in data['Solution'][1]['Variable'].items() if k.startswith('battery_E')}

之后kv 包含：

{'battery_E[b2,2,3]': 0.198967393893, 'battery_E[b1,1,1]': 0.25, 'battery_E[b1,1,2]': 0.259912707017, 'battery_E[b2,1,1]': 0.0899999972181, 'battery_E[b1,2,1]': 0.120758408109}