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 in
(task, start=True)
| 87 | |
| 88 | |
| 89 | def lazify_task(task, start=True): |
| 90 | """ |
| 91 | Given a task, remove unnecessary calls to ``list`` and ``reify``. |
| 92 | |
| 93 | This traverses tasks and small lists. We choose not to traverse down lists |
| 94 | of size >= 50 because it is unlikely that sequences this long contain other |
| 95 | sequences in practice. |
| 96 | |
| 97 | Examples |
| 98 | -------- |
| 99 | >>> def inc(x): |
| 100 | ... return x + 1 |
| 101 | >>> task = (sum, (list, (map, inc, [1, 2, 3]))) |
| 102 | >>> lazify_task(task) # doctest: +ELLIPSIS |
| 103 | (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3])) |
| 104 | """ |
| 105 | |
| 106 | if isinstance(task, GraphNode): |
| 107 | if isinstance(task, List) and len(task.args) < 50: |
| 108 | return List(*[lazify_task(arg, False) for arg in task.args]) |
| 109 | if not isinstance(task, Task): |
| 110 | return task |
| 111 | if not start and task.func in (list, reify) and isinstance(task.args[0], Task): |
| 112 | assert len(task.args) == 1 |
| 113 | task = task.args[0] |
| 114 | if task.func is _execute_subgraph: |
| 115 | subgraph, outkey, inkeys, *dependencies = task.args |
| 116 | # If there is a reify at the output of the subgraph we don't want to act |
| 117 | final_task = lazify_task(subgraph[outkey], True) |
| 118 | subgraph = { |
| 119 | k: lazify_task(v, False) for k, v in subgraph.items() if k != outkey |
| 120 | } |
| 121 | subgraph[outkey] = final_task |
| 122 | return Task( |
| 123 | task.key, |
| 124 | _execute_subgraph, |
| 125 | subgraph, |
| 126 | outkey, |
| 127 | inkeys, |
| 128 | *dependencies, |
| 129 | _data_producer=task.data_producer, |
| 130 | ) |
| 131 | return Task( |
| 132 | task.key, |
| 133 | task.func, |
| 134 | *[lazify_task(arg, False) for arg in task.args], |
| 135 | **task.kwargs, |
| 136 | ) |
| 137 | else: |
| 138 | if type(task) is list and len(task) < 50: |
| 139 | return [lazify_task(arg, False) for arg in task] |
| 140 | if not istask(task): |
| 141 | return task |
| 142 | head, tail = task[0], task[1:] |
| 143 | if not start and head in (list, reify): |
| 144 | task = task[1] |
| 145 | return lazify_task(*tail, start=False) |
| 146 | else: |