Skip to content
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.

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

fix: Use isinstance to support subclasses of values for deepmap #228

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 4 additions & 0 deletions docs/source/changelog.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
# Changelog

## 0.9.9 (XXXX-XX-XX)
- Support using subclasses of lists, tuple, and dict via deep_map. This allows usage of NamedTuple for more
semantically obvious return types.

## 0.9.8 (2024-09-06)
- Bugfix for `inputs` argument for `flow.run()`.

Expand Down
9 changes: 4 additions & 5 deletions src/pydiverse/pipedag/util/deep_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
Heavily inspired by the builtin copy module of python:
https://github.com/python/cpython/blob/main/Lib/copy.py
"""

from __future__ import annotations

from typing import Callable
Expand All @@ -19,13 +20,11 @@ def deep_map(x, fn: Callable, memo=None):
if y is not _nil:
return y

cls = type(x)

if cls == list:
if isinstance(x, list):
y = _deep_map_list(x, fn, memo)
elif cls == tuple:
elif isinstance(x, tuple):
Copy link
Member

@windiana42 windiana42 Oct 3, 2024

Choose a reason for hiding this comment

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

@btat-qc we talked about extensions like this in deep_map before. The question is what the expected behavior should be. deep_map(v, lambda x: x) should be very close to a deep copy without changing anything. With the change you proposed, the NamedTuple would be converted to a tuple if I am not mistaken.
I would suggest to explicitly handle NamedTuples and data classes here to guarantee that they are deep copied to the same type.

Copy link
Author

Choose a reason for hiding this comment

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

I see - that makes sense! Thanks I'll think about how to solve this and add some tests

y = _deep_map_tuple(x, fn, memo)
elif cls == dict:
elif isinstance(x, dict):
y = _deep_map_dict(x, fn, memo)
else:
y = fn(x)
Expand Down
8 changes: 7 additions & 1 deletion tests/test_flows/test_flow.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from __future__ import annotations
from typing import NamedTuple

import pandas as pd
import sqlalchemy as sa
Expand Down Expand Up @@ -46,9 +47,14 @@ def list_arg(x: list[pd.DataFrame]):
return Blob(x)


class Blobs(NamedTuple):
a: Blob
b: Blob


@materialize
def blob_task(x, y):
return Blob(x), Blob(y)
return Blobs(Blob(x), Blob(y))


def test_simple_flow(with_blob=True):
Expand Down
Loading