Skip to content

Instantly share code, notes, and snippets.

@autumnjolitz
Created January 7, 2023 02:23
Show Gist options
  • Save autumnjolitz/7ff233ebbac4673808e6ada5406cca0e to your computer and use it in GitHub Desktop.
Save autumnjolitz/7ff233ebbac4673808e6ada5406cca0e to your computer and use it in GitHub Desktop.
generic tagged yaml loader that can also do tuples ;) no wrapt garbage required :D
import io
from typing import Any, NamedTuple
import yaml
import yaml.nodes
class Tagged(NamedTuple):
tag: str
tag_type: str
value: Any
@classmethod
def from_yaml(cls, loader, node):
if isinstance(node, yaml.nodes.ScalarNode):
value = loader.construct_scalar(node)
tag_type = "scalar"
elif isinstance(node, yaml.nodes.SequenceNode):
tag_type = "sequence"
value = loader.construct_sequence(node)
elif isinstance(node, yaml.nodes.MappingNode):
tag_type = "mapping"
value = loader.construct_mapping(node)
else:
raise TypeError("Unknown type {node!r}")
if node.tag == 'tag:yaml.org,2002:python/tuple':
value = tuple(value)
return cls(node.tag, tag_type, value)
@classmethod
def to_yaml(cls, representer, node):
if node.tag_type == "scalar":
node = representer.represent_scalar(node.tag, node.value)
elif node.tag_type == "sequence":
node = representer.represent_sequence(node.tag, node.value)
elif node.tag_type == "mapping":
node = representer.represent_mapping(node.tag, node.value)
else:
raise TypeError("Unknown type {node!r}")
return node
class Loader(yaml.SafeLoader):
pass
Loader.add_constructor(None, Tagged.from_yaml)
class Dumper(yaml.SafeDumper):
pass
Dumper.add_representer(Tagged, Tagged.to_yaml)
def test_yml():
block = '''? !!python/tuple [0,0]
: The Hero
? !!python/tuple [0,1]
: Treasure
? !!python/tuple [1,0]
: Treasure
? !!python/tuple [1,1]
: The Dragon'''
r = yaml.load(block, Loader=Loader)
print(r)
with io.StringIO() as buf:
yaml.dump(r, buf, Dumper=Dumper)
print(buf.getvalue())
if __name__ == '__main__':
test_yml()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment