From 0f3e141b0cbf618dd7c21ac5577ca675561c0f9e Mon Sep 17 00:00:00 2001 From: Jiachen Dong Date: Tue, 15 Jul 2025 21:08:37 +0800 Subject: [PATCH 1/3] Add a missing assertion in the test. --- src/greenlet/tests/test_generator_nested.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/greenlet/tests/test_generator_nested.py b/src/greenlet/tests/test_generator_nested.py index 8d752a63..8ca17942 100644 --- a/src/greenlet/tests/test_generator_nested.py +++ b/src/greenlet/tests/test_generator_nested.py @@ -166,3 +166,4 @@ def test_nested_genlets(self): seen = [] for ii in ax(5): seen.append(ii) + self.assertEqual(seen, [1, 2, 3, 4, 5]) From 0c21634aad6040c0ef59bab3502841ba657cb536 Mon Sep 17 00:00:00 2001 From: Jiachen Dong Date: Tue, 15 Jul 2025 22:30:18 +0800 Subject: [PATCH 2/3] Add deeply nested generator tests to show that greenlet can achieve better time complexity than built-in generators. --- .../tests/test_generator_deeply_nested.py | 158 ++++++++++++++++++ .../tests/test_generator_deeply_nested2.py | 131 +++++++++++++++ 2 files changed, 289 insertions(+) create mode 100644 src/greenlet/tests/test_generator_deeply_nested.py create mode 100644 src/greenlet/tests/test_generator_deeply_nested2.py diff --git a/src/greenlet/tests/test_generator_deeply_nested.py b/src/greenlet/tests/test_generator_deeply_nested.py new file mode 100644 index 00000000..5f11c93b --- /dev/null +++ b/src/greenlet/tests/test_generator_deeply_nested.py @@ -0,0 +1,158 @@ +from greenlet import greenlet + +from . import TestCase + +def Yield(value): + """Pauses the current worker and sends a value to its parent greenlet.""" + parent = greenlet.getcurrent().parent + if not isinstance(parent, genlet): + raise RuntimeError("yield outside a genlet") + parent.switch(value) + +class _YieldFromMarker: + """Internal object that signals a `yield from` request to the trampoline.""" + def __init__(self, task): + self.task = task + +def YieldFrom(func, *args, **kwargs): + """ + Creates a marker for the trampoline to delegate to another generator. + It unwraps the decorated function to get the raw logic. + """ + # Access the original, undecorated function that the @generator stored. + raw_func = getattr(func, '_raw_func', func) + marker = _YieldFromMarker((raw_func, args, kwargs)) + Yield(marker) + +class genlet(greenlet): + """ + A greenlet that acts as a generator. It uses an internal trampoline to manage a stack of tasks, + achieving O(1) performance for each deep delegated `yield from`. + """ + def __init__(self, initial_task): + super().__init__(self.run) + self.initial_task = initial_task + self.consumer = None + + def __iter__(self): + return self + + def __next__(self): + # The consumer is the greenlet that called `next()`. + self.consumer = greenlet.getcurrent() + + # Switch to the `run` method to get the next value. + result = self.switch() + + # After the switch, the trampoline either sends a value or finishes. + if self.dead: + raise StopIteration + return result + + def run(self): + """ + The trampoline. It manages a stack of worker greenlets and never builds + a deep Python call stack itself. + """ + worker_stack = [] + + func, args, kwargs = self.initial_task + # The `active_worker` is the greenlet executing user code. Its `parent` + # is automatically set to `self` (this genlet instance) on creation. + active_worker = greenlet(func) + + # Start the first worker and capture the first value it yields. + yielded = active_worker.switch(*args, **kwargs) + + while True: + # Case 1: Delegation (`yield from`). + # The worker wants to delegate to a sub-generator. + if isinstance(yielded, _YieldFromMarker): + # Pause the current worker by pushing it onto the stack. + worker_stack.append(active_worker) + + # Create and start the new child worker. + child_func, child_args, child_kwargs = yielded.task + active_worker = greenlet(child_func) + yielded = active_worker.switch(*child_args, **child_kwargs) + continue + + # Case 2: A worker has finished. + # The worker function has returned, so its greenlet is now "dead". + if active_worker.dead: + # If there are no parent workers waiting, the whole process is done. + if not worker_stack: + break + + # A sub-generator finished. Pop its parent from the stack + # to make it the active worker again and resume it. + active_worker = worker_stack.pop() + yielded = active_worker.switch() + continue + + # Case 3: A real value was yielded. + # 1. Send the value to the consumer (the loop calling `next()`). + self.consumer.switch(yielded) + + # 2. After the consumer gets the value, control comes back here. + # Resume the active worker to ask for the next value. + yielded = active_worker.switch() + +def generator(func): + """ + Decorator that turns a function using `Yield`/`YieldFrom` into a generator. + It stores a reference to the original function to allow `YieldFrom` to work. + """ + def wrapper(*args, **kwargs): + # This wrapper is what the user calls. It creates the main genlet. + return genlet((func, args, kwargs)) + + # Store the raw function so YieldFrom can access it and bypass this wrapper. + wrapper._raw_func = func + return wrapper + + +# ============================================================================= +# Test Cases +# ============================================================================= + + +@generator +def hanoi(n, a, b, c): + if n > 1: + YieldFrom(hanoi, n - 1, a, c, b) + Yield(f'{a} -> {c}') + if n > 1: + YieldFrom(hanoi, n - 1, b, a, c) + +@generator +def make_integer_sequence(n): + if n > 1: + YieldFrom(make_integer_sequence, n - 1) + Yield(n) + +@generator +def empty_gen(): + pass + +class DeeplyNestedGeneratorTests(TestCase): + def test_hanoi(self): + results = list(hanoi(3, 'A', 'B', 'C')) + self.assertEqual( + results, + ['A -> C', 'A -> B', 'C -> B', 'A -> C', 'B -> A', 'B -> C', 'A -> C'], + ) + + def test_make_integer_sequence(self): + # It does not require `sys.setrecursionlimit` to set the recursion limit to a higher value, + # since the `yield from` logic is managed as a `task` variable on the heap, and + # the control is passed via `greenlet.switch()` instead of recursive function calls. + # + # Besides, if we use the built-in `yield` and `yield from` instead in the function + # `make_integer_sequence`, the time complexity will increase from O(n) to O(n^2). + results = list(make_integer_sequence(2000)) + self.assertEqual(results, list(range(1, 2001))) + + def test_empty_gen(self): + for _ in empty_gen(): + self.fail('empty generator should not yield anything') diff --git a/src/greenlet/tests/test_generator_deeply_nested2.py b/src/greenlet/tests/test_generator_deeply_nested2.py new file mode 100644 index 00000000..7f32faf7 --- /dev/null +++ b/src/greenlet/tests/test_generator_deeply_nested2.py @@ -0,0 +1,131 @@ +from greenlet import greenlet + +from . import TestCase + +class YieldFromMarker: + """A special object to signal a `yield from` request.""" + def __init__(self, iterable): + self.iterable = iterable + +class genlet(greenlet): + """ + A greenlet that is also an iterator, designed to wrap a function + and turn it into a generator that can be driven by a for loop. + This version includes a trampoline to handle `yield from` efficiently. + """ + def __init__(self, func, *args, **kwargs): + # We need to capture the function to run, which is stored on the class + # by the `generator` decorator. + self.func = func + self.args = args + self.kwargs = kwargs + # The stack of active iterators for the trampoline. + self.iter_stack = [] + super().__init__(self.run) + + def __iter__(self): + return self + + def __next__(self): + # Set the parent to the consumer. + self.parent = greenlet.getcurrent() # pylint:disable=attribute-defined-outside-init + # Switch to the `run` method to get the next value. + result = self.switch() + + if self: + return result + raise StopIteration + + def run(self): + """ + The trampoline. This loop drives the user's generator logic. It manages + a stack of iterators to achieve O(1) performance for each deep delegated `yield from`. + """ + # Create the top-level generator from the user's function. + top_level_generator = self.func(*self.args, **self.kwargs) + self.iter_stack.append(top_level_generator) + + while self.iter_stack: + try: + # Get the value from the top-most generator on our stack. + value = next(self.iter_stack[-1]) + + if isinstance(value, YieldFromMarker): + # It's a `yield from` request. + sub_iterable = value.iterable + # Crucially, unpack the genlet into a simple generator + # to avoid nested trampolines. + if isinstance(sub_iterable, genlet): + sub_generator = sub_iterable.func(*sub_iterable.args, **sub_iterable.kwargs) + self.iter_stack.append(sub_generator) + else: + # Support yielding from standard iterables as well, + # e.g. `yield YieldFromMarker([1, 2, 3])`. + self.iter_stack.append(iter(sub_iterable)) + else: + # It's a regular value. Pass it back to the consumer + # (which is waiting in `__next__`). + self.parent.switch(value) + + except StopIteration: + # The top-most generator is exhausted. Pop it from the stack + # and continue with the one below it. + self.iter_stack.pop() + + # If the stack is empty, the entire process is complete. + # The greenlet will die, and `__next__` will raise StopIteration. + +def generator(func): + """A decorator to create a genlet class from a function.""" + def wrapper(*args, **kwargs): + return genlet(func, *args, **kwargs) + + return wrapper + + +# ============================================================================= +# Test Cases +# ============================================================================= + + +@generator +def hanoi(n, a, b, c): + if n > 1: + yield YieldFromMarker(hanoi(n - 1, a, c, b)) + yield f'{a} -> {c}' + if n > 1: + yield YieldFromMarker(hanoi(n - 1, b, a, c)) + +@generator +def make_integer_sequence(n): + if n > 1: + yield YieldFromMarker(make_integer_sequence(n - 1)) + yield n + +@generator +def empty_gen(): + # The function body should contain at least one `yield` to make it a generator. + if False: # pylint:disable=using-constant-test + yield 1 + +class DeeplyNestedGeneratorTests(TestCase): + def test_hanoi(self): + results = list(hanoi(3, 'A', 'B', 'C')) + self.assertEqual( + results, + ['A -> C', 'A -> B', 'C -> B', 'A -> C', 'B -> A', 'B -> C', 'A -> C'], + ) + + def test_make_integer_sequence(self): + # It does not require `sys.setrecursionlimit` to set the recursion limit to a higher value, + # since the `yield from` logic is managed as a `task` variable on the heap, and + # the control is passed via `greenlet.switch()` instead of recursive function calls. + # + # Besides, if we use the built-in `yield from` instead of `yield YieldFromMarker` in the + # function `make_integer_sequence`, the time complexity will increase from O(n) to O(n^2). + results = list(make_integer_sequence(2000)) + self.assertEqual(results, list(range(1, 2001))) + + def test_empty_gen(self): + for _ in empty_gen(): + self.fail('empty generator should not yield anything') From af8d0da0aebbd1b82f3d752e4fc0d181355b3f2c Mon Sep 17 00:00:00 2001 From: Jiachen Dong Date: Tue, 15 Jul 2025 22:58:30 +0800 Subject: [PATCH 3/3] Update README to cover the example that deeply nested generators can be better than built-ins. --- README.rst | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/README.rst b/README.rst index 557488bb..f90e3194 100644 --- a/README.rst +++ b/README.rst @@ -19,12 +19,18 @@ generators; the difference with Python's own generators is that our generators can call nested functions and the nested functions can yield values too. (Additionally, you don't need a "yield" keyword. See the example in `test_generator.py -`_). +`_). +Moreover, when dealing with deeply nested generators, e.g. recursively +traversing a tree structure, due to `PEP 380 Optimizations`_ not being +implemented in CPython, our generators can achieve better time complexity +(See the example in `test_generator_deeply_nested.py +`_). Greenlets are provided as a C extension module for the regular unmodified interpreter. .. _`Stackless`: http://www.stackless.com +.. _`PEP 380 Optimizations`: https://peps.python.org/pep-0380/#optimisations Who is using Greenlet?