Repository: dask/dask
Pull request #11600 (2024-12-13 16:47:00)

Avoid constructing `NestedContainer`s in case of trivial inputs
I noticed this while debugging https://github.com/dask/distributed/issues/8959

- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  ±  0       15 suites  ±0   4h 6m 31s ⏱️ + 5m 31s
 13 275 tests +  8   12 213 ✅ +  8   1 062 💤 ±0  0 ❌ ±0 
164 721 runs  +120  141 792 ✅ +118  22 929 💤 +2  0 ❌ ±0 

Results for commit 0903b18b. ± Comparison against base commit 2eb339e0.

[test-results]:data:application/gzip;base64,H4sIAFdtXGcC/03MSw6DIBSF4a0Yxh1wLwjSzTQX1ITUR4MwMt170So6/L+TnJX1fugW9qygflRsST5e1aZA0c/T1lIbyJTXuO8CdX32a0nObYgI4sK3/2zIFRbryQ/ZeIEuhDkcEtK0XyupEY4u1xK0wYL/a0SD5rTb9d73ZzePo485GDdcWGgsAllBjdJoFfUt59pJo0EZSbUjIvb9Aa/UZiYaAQAA

hendrikmakait: @jrbourbeau: Fixed the mypy problem

Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 316f1805a..d05ee8e0f 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -84,7 +84,7 @@ _anom_count = itertools.count()
 def parse_input(obj: Any) -> object:
     """Tokenize user input into GraphNode objects
 
-    Note: This is similar to `convert_old_style_task` but does not
+    Note: This is similar to `convert_legacy_task` but does not
     - compare any values to a global set of known keys to infer references/futures
     - parse tuples and interprets them as runnable tasks
     - Deal with SubgraphCallables
@@ -101,20 +101,27 @@ def parse_input(obj: Any) -> object:
     """
     if isinstance(obj, GraphNode):
         return obj
-    if isinstance(obj, list):
-        return List(*(parse_input(o) for o in obj))
-    if isinstance(obj, set):
-        return Set(*(parse_input(o) for o in obj))
-    if isinstance(obj, tuple):
-        if is_namedtuple_instance(obj):
-            return _wrap_namedtuple_task(None, obj, parse_input)
-        return Tuple(*(parse_input(o) for o in obj))
-    if isinstance(obj, dict):
-        return Dict({k: parse_input(v) for k, v in obj.items()})
 
     if isinstance(obj, TaskRef):
         return Alias(obj.key)
 
+    if isinstance(obj, dict):
+        parsed_dict = {k: parse_input(v) for k, v in obj.items()}
+        if any(isinstance(v, GraphNode) for v in parsed_dict.values()):
+            return Dict(parsed_dict)
+
+    if isinstance(obj, (list, set, tuple)):
+        parsed_collection = tuple(parse_input(o) for o in obj)
+        if any(isinstance(o, GraphNode) for o in parsed_collection):
+            if isinstance(obj, list):
+                return List(*parsed_collection)
+            if isinstance(obj, set):
+                return Set(*parsed_collection)
+            if isinstance(obj, tuple):
+                if is_namedtuple_instance(obj):
+                    return _wrap_namedtuple_task(None, obj, parse_input)
+                return Tuple(*parsed_collection)
+
     return obj
 
 


Test patch:
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 60885ec1f..c62447643 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -538,6 +538,11 @@ def test_pickle_literals():
     assert roundtripped == obj
 
 
+@pytest.mark.parametrize("obj", [set, {0}, [], [1], {}, {2: 3}, (), (4,)])
+def test_parse_non_task_inputs(obj):
+    assert parse_input(obj) == obj
+
+
 def test_resolve_aliases():
     tasks = [
         Alias("bar", "foo"),