diff --git a/src/sciline/scheduler.py b/src/sciline/scheduler.py index 1da9a05e..d13f462b 100644 --- a/src/sciline/scheduler.py +++ b/src/sciline/scheduler.py @@ -61,11 +61,12 @@ def get( raise CycleError from e dependents = _count_dependents(dependencies) + requested = set(keys) results: dict[Hashable, Any] = {} with reporter.run_computation(graph.values()): for t in tasks: results[t] = reporter.call_provider_with_reporting(graph[t], results) - _consume_arguments(graph[t], dependents, results) + _consume_arguments(graph[t], dependents, results, requested) return tuple(results[key] for key in keys) @@ -81,11 +82,19 @@ def _count_dependents(dependencies: dict[type, tuple[type, ...]]) -> Counter[typ def _consume_arguments( - provider: Provider, counts: Counter[type], results: dict[Hashable, object] + provider: Provider, + counts: Counter[type], + results: dict[Hashable, object], + requested: set[Hashable], ) -> None: + """Discard results that no remaining provider needs. + + Requested keys are kept: they are returned to the caller, so their consumer + count reaching zero does not mean they are no longer needed. + """ for arg in provider.arg_spec.keys(): counts[arg] -= 1 - if counts[arg] == 0: + if counts[arg] == 0 and arg not in requested: del results[arg] diff --git a/tests/pipeline_test.py b/tests/pipeline_test.py index 00bd72b2..11f971cb 100644 --- a/tests/pipeline_test.py +++ b/tests/pipeline_test.py @@ -879,6 +879,20 @@ def test_compute_with_NaiveScheduler() -> None: assert res == 1.5 +def test_compute_returns_requested_key_that_another_requested_key_depends_on( + scheduler: sl.scheduler.Scheduler, +) -> None: + def make_int_local() -> int: + return 3 + + def int_to_float_local(x: int) -> float: + return 0.5 * x + + pipeline = sl.Pipeline([int_to_float_local, make_int_local]) + # int is both a requested key and an argument of the provider of float. + assert pipeline.compute((int, float), scheduler=scheduler) == {int: 3, float: 1.5} + + def test_bind_and_call_no_function() -> None: pipeline = sl.Pipeline([make_int]) assert pipeline.bind_and_call(()) == ()