Skip to content

Handle Uuid FullLoader with yaml de/serializer #310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
unkcpz merged 3 commits into from
Feb 21, 2025
Merged

Handle Uuid FullLoader with yaml de/serializer #310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b64dcf9
Handle Uuid FullLoader with yaml de/serializer
unkcpz Jan 30, 2025
92d6d8a
Merge branch 'master' into handle-uuid-yml-in-message-respond
unkcpz Feb 18, 2025
3165b2c
Explicit argument pass
unkcpz Feb 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/plumpy/persistence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import inspect
import os
import pickle
import uuid
from types import MethodType
from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, Iterable, List, Optional, Set, TypeVar, Union

Expand Down Expand Up @@ -35,6 +36,19 @@
from .processes import Process


def uuid_representer(dumper, data): # type: ignore
return dumper.represent_scalar('!uuid', str(data))


def uuid_constructor(loader, node): # type: ignore
value = loader.construct_scalar(node)
return uuid.UUID(value)


yaml.add_representer(uuid.UUID, uuid_representer)
yaml.add_constructor('!uuid', uuid_constructor)


class Bundle(dict):
def __init__(self, savable: 'Savable', save_context: Optional['LoadSaveContext'] = None, dereference: bool = False):
"""
Expand Down
5 changes: 4 additions & 1 deletion tests/rmq/test_communicator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,16 @@ def loop_communicator():
message_exchange = f'{__file__}.{shortuuid.uuid()}'
task_exchange = f'{__file__}.{shortuuid.uuid()}'
task_queue = f'{__file__}.{shortuuid.uuid()}'
encoder = functools.partial(yaml.dump, encoding='utf-8')
decoder = functools.partial(yaml.load, Loader=yaml.FullLoader)

thread_communicator = rmq.RmqThreadCommunicator.connect(
connection_params={'url': 'amqp://guest:guest@localhost:5672/'},
message_exchange=message_exchange,
task_exchange=task_exchange,
task_queue=task_queue,
decoder=functools.partial(yaml.load, Loader=yaml.Loader),
Copy link
Member Author

@unkcpz unkcpz Feb 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks! One thing is not clear to me: if the default loader in the plugin is the safe one, why is there no error before this PR. The safe loader will raise an error if the uuid is not serializable.

@superstar54 This line matters. If I didn't add the uuid_representer and uuid_constructor, after remove this unsafeloader decoder. The test hang because it can not deserialize the message.

Copy link
Member Author

@unkcpz unkcpz Feb 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@superstar54 I guess your previous comment was gonna to reply this one?

I updated, by passing explicitly the decoder and encoder. Here actually it is a test, not the real code path used.

Copy link
Member

@superstar54 superstar54 Feb 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's still not clear to me. I am guessing you are adding the uuid_representer and uuid_constructor for safe load. But where is the test for them?

Copy link
Member Author

@unkcpz unkcpz Feb 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

uuid_representer and uuid_constructor

If I didn't add these and I remove the decoder line in loop_communicator, the test will fail (hang actually, because RMQ never get the correct message).

Copy link
Member

@superstar54 superstar54 Feb 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since you are adding a constructor, and your goal is to have the safe load, as suggested in my comment, it would be good to use yaml.SafeLoader instead of yaml.FullLoader, unless there is still another message type that also needs a constructor.

If there is another message type, you can add a constructor for it, as you already did for the uuid. It would be good if you could make a list of the possible message types.

If you think yaml.FullLoader is already safe enough, I am also fine.

Copy link
Member Author

@unkcpz unkcpz Feb 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you think yaml.FullLoader is already safe enough, I am also fine.

The scope if this PR is to revert the original unnecessary change and using the default loader of kiwipy.
The plan for the message passing part is to use msgpack to replace. So for this one I'll not changing the test above.

encoder=encoder,
decoder=decoder,
)

loop = asyncio.get_event_loop()
Expand Down
2 changes: 1 addition & 1 deletion tests/test_persistence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,6 @@ def test_bundle_yaml(self):
bundle = plumpy.Bundle(Save1())
represent = yaml.dump({'bundle': bundle})

bundle_loaded = yaml.load(represent, Loader=yaml.Loader)['bundle']
bundle_loaded = yaml.load(represent, Loader=yaml.UnsafeLoader)['bundle']
self.assertIsInstance(bundle_loaded, plumpy.Bundle)
self.assertDictEqual(bundle_loaded, Save1().save())