Repository: dask/dask
File: dask/bag/core.py (Python)

Code before:
    Given a task, remove unnecessary calls to ``list`` and ``reify``.

    This traverses tasks and small lists.  We choose not to traverse down lists
    of size >= 50 because it is unlikely that sequences this long contain other
    sequences in practice.

    Examples
    --------
    >>> def inc(x):
    ...     return x + 1
    >>> task = (sum, (list, (map, inc, [1, 2, 3])))
    >>> lazify_task(task)  # doctest: +ELLIPSIS
    (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3]))
    """
    if isinstance(task, GraphNode):
        if isinstance(task, List) and len(task.args) < 50:
            return List(*[lazify_task(arg, False) for arg in task.args])
        if not isinstance(task, Task):
            return task
        if not start and task.func in (list, reify) and isinstance(task.args[0], Task):
            assert len(task.args) == 1
            task = task.args[0]
        if task.func is _execute_subgraph:
            subgraph = task.args[0]
            outkey = task.args[1]
            # If there is a reify at the start of this we don't want to act
            final_task = lazify_task(subgraph[outkey], True)
            subgraph = {
                k: lazify_task(v, False) for k, v in subgraph.items() if k != outkey
            }
            subgraph[outkey] = final_task
            return Task(
                task.key,
                _execute_subgraph,
                subgraph,
                outkey,
                *task.args[2:],
                **task.kwargs,
            )
        return Task(
            task.key,
            task.func,
            *[lazify_task(arg, False) for arg in task.args],
            **task.kwargs,
        )
    else:
        if type(task) is list and len(task) < 50:
            return [lazify_task(arg, False) for arg in task]
        if not istask(task):
            return task
        head, tail = task[0], task[1:]

Review comment (suggestion):
```suggestion
            # If there is a reify at the output of the subgraph we don't want to act
```
I think this is what you mean and might be a bit clearer.

Code after:
    Given a task, remove unnecessary calls to ``list`` and ``reify``.

    This traverses tasks and small lists.  We choose not to traverse down lists
    of size >= 50 because it is unlikely that sequences this long contain other
    sequences in practice.

    Examples
    --------
    >>> def inc(x):
    ...     return x + 1
    >>> task = (sum, (list, (map, inc, [1, 2, 3])))
    >>> lazify_task(task)  # doctest: +ELLIPSIS
    (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3]))
    """
    if isinstance(task, GraphNode):
        if isinstance(task, List) and len(task.args) < 50:
            return List(*[lazify_task(arg, False) for arg in task.args])
        if not isinstance(task, Task):
            return task
        if not start and task.func in (list, reify) and isinstance(task.args[0], Task):
            assert len(task.args) == 1
            task = task.args[0]
        if task.func is _execute_subgraph:
            subgraph = task.args[0]
            outkey = task.args[1]
            # If there is a reify at the output of the subgraph we don't want to act
            final_task = lazify_task(subgraph[outkey], True)
            subgraph = {
                k: lazify_task(v, False) for k, v in subgraph.items() if k != outkey
            }
            subgraph[outkey] = final_task
            return Task(
                task.key,
                _execute_subgraph,
                subgraph,
                outkey,
                *task.args[2:],
                **task.kwargs,
            )
        return Task(
            task.key,
            task.func,
            *[lazify_task(arg, False) for arg in task.args],
            **task.kwargs,
        )
    else:
        if type(task) is list and len(task) < 50:
            return [lazify_task(arg, False) for arg in task]
        if not istask(task):
            return task
        head, tail = task[0], task[1:]