Extract collections in preparation for compute/persist/etc... Intended use is to find all collections in a set of (possibly nested) python objects, do something to them (compute, etc...), then repackage them in equivalent python objects. Parameters ---------- *args
(*args, traverse=True)
| 513 | |
| 514 | |
| 515 | def unpack_collections(*args, traverse=True): |
| 516 | """Extract collections in preparation for compute/persist/etc... |
| 517 | |
| 518 | Intended use is to find all collections in a set of (possibly nested) |
| 519 | python objects, do something to them (compute, etc...), then repackage them |
| 520 | in equivalent python objects. |
| 521 | |
| 522 | Parameters |
| 523 | ---------- |
| 524 | *args |
| 525 | Any number of objects. If it is a dask collection, it's extracted and |
| 526 | added to the list of collections returned. By default, python builtin |
| 527 | collections are also traversed to look for dask collections (for more |
| 528 | information see the ``traverse`` keyword). |
| 529 | traverse : bool, optional |
| 530 | If True (default), builtin python collections are traversed looking for |
| 531 | any dask collections they might contain. |
| 532 | |
| 533 | Returns |
| 534 | ------- |
| 535 | collections : list |
| 536 | A list of all dask collections contained in ``args`` |
| 537 | repack : callable |
| 538 | A function to call on the transformed collections to repackage them as |
| 539 | they were in the original ``args``. |
| 540 | """ |
| 541 | |
| 542 | collections = [] |
| 543 | repack_dsk = {} |
| 544 | |
| 545 | collections_token = uuid.uuid4().hex |
| 546 | |
| 547 | def _unpack(expr): |
| 548 | if is_dask_collection(expr): |
| 549 | tok = tokenize(expr) |
| 550 | if tok not in repack_dsk: |
| 551 | repack_dsk[tok] = Task( |
| 552 | tok, getitem, TaskRef(collections_token), len(collections) |
| 553 | ) |
| 554 | collections.append(expr) |
| 555 | return TaskRef(tok) |
| 556 | |
| 557 | tok = uuid.uuid4().hex |
| 558 | tsk: DataNode | Task # type: ignore[annotation-unchecked] |
| 559 | if not traverse: |
| 560 | tsk = DataNode(None, expr) |
| 561 | else: |
| 562 | # Treat iterators like lists |
| 563 | typ = list if isinstance(expr, Iterator) else type(expr) |
| 564 | if typ in (list, tuple, set): |
| 565 | tsk = Task(tok, typ, List(*[_unpack(i) for i in expr])) |
| 566 | elif typ in (dict, OrderedDict): |
| 567 | tsk = Task( |
| 568 | tok, typ, Dict({_unpack(k): _unpack(v) for k, v in expr.items()}) |
| 569 | ) |
| 570 | elif dataclasses.is_dataclass(expr) and not isinstance(expr, type): |
| 571 | tsk = Task( |
| 572 | tok, |