import yaml
yaml.add_constructor("!flatten", construct_flat_list)
def flatten_sequence(sequence: yaml.Node) -> Iterator[str]:
"""Flatten a nested sequence to a list of strings
A nested structure is always a SequenceNode
"""
if isinstance(sequence, yaml.ScalarNode):
yield sequence.value
return
if not isinstance(sequence, yaml.SequenceNode):
raise TypeError(f"'!flatten' can only flatten sequence nodes, not {sequence}")
for el in sequence.value:
if isinstance(el, yaml.SequenceNode):
yield from flatten_sequence(el)
elif isinstance(el, yaml.ScalarNode):
yield el.value
else:
raise TypeError(f"'!flatten' can only take scalar nodes, not {el}")
def construct_flat_list(loader: yaml.Loader, node: yaml.Node) -> List[str]:
"""Make a flat list, should be used with '!flatten'
Args:
loader: Unused, but necessary to pass to `yaml.add_constructor`
node: The passed node to flatten
"""
return list(flatten_sequence(node))