Skip to content

Commit

Permalink
feat: Extend YAML serialization to allow Python tuples (#7853)
Browse files Browse the repository at this point in the history
  • Loading branch information
shadeMe authored Jun 14, 2024
1 parent 24518df commit 2f6ae41
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 1 deletion.
12 changes: 11 additions & 1 deletion haystack/marshal/yaml.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,21 @@
import yaml


# Custom YAML safe loader that supports loading Python tuples
class YamlLoader(yaml.SafeLoader): # pylint: disable=too-many-ancestors
def construct_python_tuple(self, node: yaml.SequenceNode):
"""Construct a Python tuple from the sequence."""
return tuple(self.construct_sequence(node))


YamlLoader.add_constructor("tag:yaml.org,2002:python/tuple", YamlLoader.construct_python_tuple)


class YamlMarshaller:
def marshal(self, dict_: Dict[str, Any]) -> str:
"""Return a YAML representation of the given dictionary."""
return yaml.dump(dict_)

def unmarshal(self, data_: Union[str, bytes, bytearray]) -> Dict[str, Any]:
"""Return a dictionary from the given YAML data."""
return yaml.safe_load(data_)
return yaml.load(data_, Loader=YamlLoader)
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
enhancements:
- |
Pipeline serialization to YAML now supports tuples as field values.
3 changes: 3 additions & 0 deletions test/marshal/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
42 changes: 42 additions & 0 deletions test/marshal/test_yaml.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
import pytest

from haystack.marshal.yaml import YamlMarshaller


@pytest.fixture
def yaml_data():
return {"key": "value", 1: 0.221, "list": [1, 2, 3], "tuple": (1, None, True), "dict": {"set": {False}}}


@pytest.fixture
def serialized_yaml_str():
return """key: value
1: 0.221
list:
- 1
- 2
- 3
tuple: !!python/tuple
- 1
- null
- true
dict:
set: !!set
false: null
"""


def test_yaml_marshal(yaml_data, serialized_yaml_str):
marshaller = YamlMarshaller()
marshalled = marshaller.marshal(yaml_data)
assert isinstance(marshalled, str)
assert marshalled.strip().replace("\n", "") == serialized_yaml_str.strip().replace("\n", "")


def test_yaml_unmarshal(yaml_data, serialized_yaml_str):
marshaller = YamlMarshaller()
unmarshalled = marshaller.unmarshal(serialized_yaml_str)
assert unmarshalled == yaml_data

0 comments on commit 2f6ae41

Please sign in to comment.