From 904ef7796ddddfafec3eb0b3c5d30789c2449f3f Mon Sep 17 00:00:00 2001 From: Elias Strauss Date: Wed, 5 Aug 2026 15:02:07 +0200 Subject: [PATCH] Support tuple unpacking of DataOps `a, b, c = data_op` now creates a node per target instead of raising an error, so a deferred function returning several values can be unpacked as usual. Python calls `iter()` on the right-hand side from an `UNPACK_SEQUENCE` instruction whose argument is the number of targets, so `DataOp.__iter__` reads that count off the calling frame's bytecode and creates an `AsTuple` node (which converts the result to a tuple and checks its length) plus one `GetItem` per target. A wrong number of targets is reported eagerly when a preview value is available, and at evaluation time otherwise. Only `UNPACK_SEQUENCE` is recognised, so anything else (`for` loops, `list(data_op)`, `f(*data_op)`, `a, *rest = data_op`) keeps raising the previous TypeError, as does any case where the bytecode cannot be inspected. Indexing into the result stays available as a portable alternative. --- CHANGES.rst | 4 + doc/modules/data_ops/basics/control_flow.rst | 35 +++++---- skrub/_data_ops/_data_ops.py | 57 ++++++++++++++ skrub/_data_ops/tests/test_data_ops.py | 27 +++++++ skrub/_data_ops/tests/test_errors.py | 79 ++++++++++++++++++++ 5 files changed, 187 insertions(+), 15 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 139c6b995..b53f5914d 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -12,6 +12,10 @@ Ongoing development New Features ------------ +- It is now possible to unpack a :class:`DataOp` that evaluates to several + values, for example ``train, test = data_op``. Each target becomes a DataOp + that extracts one of the values. + :pr:`2243` by :user:`Elias Strauss `. Changes ------- diff --git a/doc/modules/data_ops/basics/control_flow.rst b/doc/modules/data_ops/basics/control_flow.rst index 0e5dd4782..678c9d8b1 100644 --- a/doc/modules/data_ops/basics/control_flow.rst +++ b/doc/modules/data_ops/basics/control_flow.rst @@ -106,25 +106,30 @@ Result: Unpacking multiple outputs from deferred functions -------------------------------------------------- -When a deferred function returns more than one value, you cannot unpack the -result directly because unpacking iterates over the result. Iteration is not -supported on DataOps until evaluation. - -In general, it is recommended that deferred functions return a single -value whenever possible. Returning multiple outputs should be avoided unless -strictly necessary, as it makes downstream usage more complex. - -Instead, keep the result as a single DataOp and index into it: +When a deferred function returns more than one value, we can unpack the result +as usual: >>> test = skrub.var("test", [1, 2]) >>> @skrub.deferred ... def process_test_data(test): -... left = test[0] -... right = test[1] -... return left, right ->>> res = test.skb.apply_func(process_test_data) ->>> left = res[0] ->>> right = res[1] +... return test[0], test[1] +>>> left, right = process_test_data(test) +>>> left + +Result: +――――――― +1 + +Each target becomes a DataOp that extracts one of the values, once the +computation runs. To do so, skrub needs to know how many values we are +unpacking: it finds out by inspecting the bytecode of the assignment. This +works on CPython, but it is not guaranteed by the language, and it does not +cover starred targets such as ``first, *rest = ...`` (the number of values is +not known in that case). When in doubt, we can always keep the result as a +single DataOp and index into it, which is equivalent and always works: + +>>> res = process_test_data(test) +>>> left, right = res[0], res[1] :func:`deferred` is useful not only for our own functions, but also when we need to call module-level functions from a library. For example, to delay the diff --git a/skrub/_data_ops/_data_ops.py b/skrub/_data_ops/_data_ops.py index 61ad4b1e4..4d26e9214 100644 --- a/skrub/_data_ops/_data_ops.py +++ b/skrub/_data_ops/_data_ops.py @@ -34,6 +34,7 @@ import operator import pathlib import re +import sys import textwrap import traceback import types @@ -197,6 +198,28 @@ def _format_data_op_creation_stack(): return traceback.format_list(stack) +def _unpack_arity(): + """Number of targets in the unpacking assignment being executed, if any. + + Read from the caller's ``UNPACK_SEQUENCE`` instruction, or None if not found. + """ + if (getframe := getattr(sys, "_getframe", None)) is None: + return None + try: + # skip the frames of this function and of DataOp.__iter__ + frame = getframe(2) + for instruction in dis.get_instructions(frame.f_code): + if instruction.offset == frame.f_lasti: + if instruction.opname == "UNPACK_SEQUENCE": + return instruction.arg + return None + except Exception: + # this is best-effort introspection: anything unexpected must fall back + # on refusing to iterate, not raise something else. + pass + return None + + class DataOpImpl: """Base class for all kinds of DataOps (computation graph nodes). @@ -700,6 +723,13 @@ def __bool__(self): ) def __iter__(self): + # Unpacking (`a, b = data_op`) is supported: we know how many values are + # expected, so we can create a node for each of them. Any other kind of + # iteration would need the length of the result, which is unknown until + # the DataOp is evaluated. + if (arity := _unpack_arity()) is not None: + values = unpack(self, arity) + return iter([values[i] for i in range(arity)]) raise TypeError( "This object is a DataOp that will be evaluated later, " "when your learner runs. So it is not possible to eagerly " @@ -1662,6 +1692,33 @@ def pretty_repr(self): return f"[{_get_preview(self.key)!r}]" +class AsTuple(DataOpImpl): + """Node created by unpacking a DataOp, e.g. ``a, b = data_op``.""" + + _fields = ["iterable", "expected_length"] + + def compute(self, e, mode, environment): + # converting to a tuple (rather than indexing into the result directly) + # allows unpacking any iterable, and makes sure an iterator is consumed + # only once even though each target indexes into this node. + result = tuple(e.iterable) + expected, got = e.expected_length, len(result) + if got != expected: + problem = "not enough" if got < expected else "too many" + raise ValueError( + f"{problem} values to unpack (expected {expected}, got {got})" + ) + return result + + def __repr__(self): + return f"<{self.__class__.__name__} {short_repr(self.iterable)}>" + + +@checked_data_op_constructor +def unpack(iterable, expected_length): + return DataOp(AsTuple(iterable, expected_length)) + + class Call(DataOpImpl): _fields = [ "func", diff --git a/skrub/_data_ops/tests/test_data_ops.py b/skrub/_data_ops/tests/test_data_ops.py index 23a926182..fa3ea7cb3 100644 --- a/skrub/_data_ops/tests/test_data_ops.py +++ b/skrub/_data_ops/tests/test_data_ops.py @@ -46,6 +46,33 @@ def test_slice(): assert c.skb.eval({"a": list(range(10, 20)), "b": 5}) == [10, 11, 12, 13, 14] +def test_unpacking(): + a = skrub.var("a", [10, 20, 30]) + first, second, third = a + assert repr(first).startswith("") + assert (first.skb.eval(), second.skb.eval(), third.skb.eval()) == (10, 20, 30) + assert (first + second).skb.eval({"a": [1, 2, 3]}) == 3 + + +def test_unpacking_iterable(): + # the result does not need to be a sequence, and an iterator must be + # consumed only once even though each target indexes into it. + @skrub.deferred + def f(): + yield 1 + yield 2 + yield 3 + + a, b, c = f() + assert (a + b + c).skb.eval() == 6 + + +def test_unpacking_nested(): + a = skrub.var("a", [1, [2, 3]]) + first, (second, third) = a + assert (first.skb.eval(), second.skb.eval(), third.skb.eval()) == (1, 2, 3) + + def test_estimator(): c = skrub.var("c") e = skrub.as_data_op(LogisticRegression(C=c)) diff --git a/skrub/_data_ops/tests/test_errors.py b/skrub/_data_ops/tests/test_errors.py index 2d17c00ec..3f9d90c0b 100644 --- a/skrub/_data_ops/tests/test_errors.py +++ b/skrub/_data_ops/tests/test_errors.py @@ -2,6 +2,7 @@ import re import sys import traceback +import types import numpy as np import pytest @@ -11,6 +12,7 @@ from sklearn.preprocessing import StandardScaler import skrub +from skrub._data_ops import _data_ops from skrub._utils import PassThrough from skrub.conftest import skip_polars_installed_without_pyarrow @@ -28,6 +30,83 @@ def test_for(): pass +def test_star_unpacking(): + # unlike `a, b = data_op`, iteration with a starred target does not tell us + # how many values are expected so it remains unsupported. + a = skrub.var("a", [1, 2, 3]) + with pytest.raises( + TypeError, match=".*it is not possible to eagerly iterate over it" + ): + _first, *_rest = a + with pytest.raises( + TypeError, match=".*it is not possible to eagerly iterate over it" + ): + (lambda *args: None)(*a) + + +def test_unpacking_wrong_number_of_targets(): + a = skrub.var("a", [1, 2, 3]) + with pytest.raises( + RuntimeError, + match=r"(?s)Evaluation of 'unpack\(\)' failed" + r".*too many values to unpack \(expected 2, got 3\)", + ): + _first, _second = a + with pytest.raises( + RuntimeError, + match=r"(?s)Evaluation of 'unpack\(\)' failed" + r".*not enough values to unpack \(expected 4, got 3\)", + ): + _first, _second, _third, _fourth = a + + +def _bytecode_inspection_failure(code): + raise RuntimeError("cannot inspect bytecode") + + +@pytest.mark.parametrize( + "module_name, replacement", + [ + # a Python implementation that does not provide sys._getframe + pytest.param("sys", types.SimpleNamespace(), id="no_getframe"), + # inspecting the bytecode fails + pytest.param( + "dis", + types.SimpleNamespace(get_instructions=_bytecode_inspection_failure), + id="inspection_error", + ), + # the instruction being executed is not found in the bytecode + pytest.param( + "dis", + types.SimpleNamespace(get_instructions=lambda code: iter(())), + id="instruction_not_found", + ), + ], +) +def test_unpacking_without_bytecode_inspection(monkeypatch, module_name, replacement): + # `a, b = data_op` relies on finding the UNPACK_SEQUENCE instruction that is + # being executed. When that is not possible we fall back on refusing to + # iterate, as users can always index into the result instead. + monkeypatch.setattr(_data_ops, module_name, replacement) + a = skrub.var("a", [1, 2]) + with pytest.raises( + TypeError, match=".*it is not possible to eagerly iterate over it" + ): + _first, _second = a + + +def test_unpacking_wrong_number_of_targets_at_runtime(): + # without a value for 'a' the length is only known when the plan runs + a = skrub.var("a") + first, _second = a + if sys.version_info < (3, 11): + err_t, err_msg = RuntimeError, "Evaluation of node