Repository: dask/dask
File: dask/_task_spec.py (Python)

Code before:
class List(NestedContainer):
    constructor = klass = list


class Tuple(NestedContainer):
    constructor = klass = tuple


class Set(NestedContainer):
    constructor = klass = set


class Dict(NestedContainer, Mapping):
    klass = dict

    def __init__(self, /, *args: Any, **kwargs: Any):
        if args:
            assert not kwargs
            if len(args) == 1:
                if isinstance(args[0], dict):
                    kwargs = args[0]
                    args = ()
                elif isinstance(args[0], (list, tuple)):
                    from dask.core import flatten

                    args = tuple(flatten(args[0]))
                    if len(args) % 2 != 0:
                        raise ValueError("Invalid number of arguments provided")
            else:
                raise ValueError(
                    "Invalid number of arguments provided. "
                    "Either provide a single dict or a list of key-value pairs"
                )
        if kwargs:
            assert not args
            args = tuple(itertools.chain(*kwargs.items()))

        super().__init__(*args)

    def __repr__(self):
        values = ", ".join(f"{k}: {v}" for k, v in batched(self.args, 2, strict=True))
        return f"Dict({values})"

    def substitute(
        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
    ) -> Dict:
        subs_filtered = {
            k: v for k, v in subs.items() if k in self.dependencies and k != v
        }
        if not subs_filtered:
            return self

Review comment (suggestion):
Looks like `flatten` is recursive and goes too far. We end up transforming

```python
([['columns', ['a', 'b']]],)
```

to

```
('columns', 'a', 'b')
```

I think we're wanting `("columns", ["a", "b"])` here instead.

This is coming from kwargs like `{"columns": ["a", b"]}` which leads to

```
(Pdb) pp keyargs
['columns']
(Pdb) pp valargs
[['a', 'b']]
```

There's probably a better fix, but we can handle this special case with

```diff
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 92d10b52e..02d9e84e2 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -915,7 +915,15 @@ class Dict(NestedContainer, Mapping):
                 elif isinstance(args[0], (list, tuple)):
                     from dask.core import flatten
 
-                    args = tuple(flatten(args[0]))
+                    if (
+                        len(args[0]) == 1
+                        and isinstance(args[0][0], (list, tuple))
+                        and len(args[0][0]) == 2
+                    ):
+                        args = tuple(args[0][0])
+                    else:
+                        args = tuple(flatten(args[0]))
+
                     if len(args) % 2 != 0:
                         raise ValueError("Invalid number of arguments provided")
             else:
```

Code after:
class List(NestedContainer):
    constructor = klass = list


class Tuple(NestedContainer):
    constructor = klass = tuple


class Set(NestedContainer):
    constructor = klass = set


class Dict(NestedContainer, Mapping):
    klass = dict

    def __init__(self, /, *args: Any, **kwargs: Any):
        if args:
            assert not kwargs
            if len(args) == 1:
                args = args[0]
                if isinstance(args, dict):  # type: ignore
                    args = tuple(itertools.chain(*args.items()))  # type: ignore
                elif isinstance(args, (list, tuple)):
                    if all(
                        len(el) == 2 if isinstance(el, (list, tuple)) else False
                        for el in args
                    ):
                        args = tuple(itertools.chain(*args))
                else:
                    raise ValueError("Invalid argument provided")

            if len(args) % 2 != 0:
                raise ValueError("Invalid number of arguments provided")

        elif kwargs:
            assert not args
            args = tuple(itertools.chain(*kwargs.items()))

        super().__init__(*args)

    def __repr__(self):
        values = ", ".join(f"{k}: {v}" for k, v in batched(self.args, 2, strict=True))
        return f"Dict({values})"

    def substitute(
        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
    ) -> Dict:
        subs_filtered = {
            k: v for k, v in subs.items() if k in self.dependencies and k != v
        }
        if not subs_filtered: