Created
January 7, 2023 02:23
-
-
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
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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