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