Repository: dask/dask
File: dask/tests/test_task_spec.py (Python)

Code before:
        == 126
    )


def apply(func, args, kwargs=None):
    return func(*args, **(kwargs or {}))


def test_redirect_to_subgraph_callable():
    subgraph = SubgraphCallable(
        {
            "out": (apply, func3, ["_0"], (dict, [["some", "dict"]])),
        },
        "out",
        ("_0",),
    )
    dsk = {
        "alias": "foo",
        "foo": (func, (subgraph, (func2, "bar", "baz")), "second_arg"),
    }
    new_dsk = convert_and_verify_keys(dsk)

    t = new_dsk["alias"]
    assert t.dependencies == {"foo"}
    t = new_dsk["foo"]
    assert not t.dependencies

    assert t() == "bar=baz//some=dict-second_arg"


def test_inline():
    t = Task("key-1", func, "a", TaskRef("b"))

    assert t.inline({"b": DataNode(None, "b")})() == "a-b"

    assert t.inline({"b": Task(None, func, "b", "c")})() == "a-b-c"

    t2 = Task("key-1", func, TaskRef("a"), TaskRef("b"))
    inline2 = t2.inline({"a": DataNode(None, "a")})
    assert inline2.dependencies == {"b"}

    assert inline2({"b": "foo"}) == "a-foo"

    t = Task("foo", func, t, t2)
    assert t.dependencies == {"b", "a"}
    inlined = t.inline({"a": DataNode(None, "foo")})
    assert inlined.dependencies == {"b"}
    assert inlined({"b": "bar"}) == "a-bar-foo-bar"

    t = Task("key-2", func, t, t2)
    t = Task("key-2", func, {"foo": t, "bar": t2})

Review comment (suggestion):
This test shows probably best what the impact of this change will be. Passing nested objects, whether they are references or actual tasks is no longer supported and will silently produce odd results.

We may want to implement a test-only validation step that still recurses into data structures to ensure we're not messing this up accidentally but ultimately this will safe a lot of computation time during runtime.

i.e. if I genuinely want to pass a list, this would have to look like...

```python
In [1]: from dask._task_spec import TaskRef, Task

In [2]: def func(a):
   ...:     return "-".join(a)
   ...:
   ...: def aslist(*args):
   ...:     return list(args)
   ...:
   ...: task_with_nested_refs = Task("key", func, Task(None, aslist, TaskRef("foo"), TaskRef("bar")))

In [3]: task_with_nested_refs.dependencies
Out[3]: frozenset({'bar', 'foo'})

In [4]: task_with_nested_refs({"foo": "foo", "bar": "bar"})
Out[4]: 'foo-bar'
```

This is why I had the SequenceOfTasks / DictOfTasks classes initially (though they were also poorly implemented. I think it makes sense to add dedicated functions/classes for these operations to avoid having to construct the above stuff manually all the time

cc @hendrikmakait 

Code after:
def test_ensure_slots(inst):
    assert not hasattr(inst, "__dict__")


class PlainNamedTuple(namedtuple("PlainNamedTuple", "value")):
    """Namedtuple with a default constructor."""


class NewArgsNamedTuple(namedtuple("NewArgsNamedTuple", "ab, c")):
    """Namedtuple with a custom constructor."""

    def __new__(cls, a, b, c):
        return super().__new__(cls, f"{a}-{b}", c)

    def __getnewargs__(self):
        return *self.ab.split("-"), self.c


class NewArgsExNamedTuple(namedtuple("NewArgsExNamedTuple", "ab, c, k, v")):
    """Namedtuple with a custom constructor including keywords-only arguments."""

    def __new__(cls, a, b, c, **kw):
        return super().__new__(cls, f"{a}-{b}", c, tuple(kw.keys()), tuple(kw.values()))

    def __getnewargs_ex__(self):
        return (*self.ab.split("-"), self.c), dict(zip(self.k, self.v))


@pytest.mark.parametrize(
    "typ, args, kwargs",
    [
        (PlainNamedTuple, ["some-data"], {}),
        (NewArgsNamedTuple, ["some", "data", "more"], {}),
        (NewArgsExNamedTuple, ["some", "data", "more"], {"another": "data"}),
    ],
)
def test_parse_graph_namedtuple_legacy(typ, args, kwargs):
    def func(x):
        return x

    dsk = {"foo": (func, typ(*args, **kwargs))}
    new_dsk = convert_and_verify_keys(dsk)

    assert new_dsk["foo"]() == typ(*args, **kwargs)


@pytest.mark.parametrize(
    "typ, args, kwargs",
    [
        (PlainNamedTuple, ["some-data"], {}),
        (NewArgsNamedTuple, ["some", "data", "more"], {}),