"""Support for tail calls in Python, including context managers. Do not use this code in production. """ import functools import inspect import sys # pylint: disable=too-few-public-methods class _ReifiedCall: def __init__(self, function, args, kwargs): self.function = function self.args = args self.kwargs = kwargs class _Tail(BaseException): def __init__(self, reified_call, *args, **kwargs): super().__init__(reified_call, *args, **kwargs) self.reified_call = reified_call self.exits = [] class _DeferExits: def __init__(self): self.exits = [] def __enter__(self): return self.exits def __exit__(self, exc_type, exc_value, traceback): for exit_group in reversed(self.exits): for exit_ in exit_group: if exit_(exc_type, exc_value, traceback): exc_type = exc_value = traceback = None return exc_type is None class tail_call: # pylint: disable=invalid-name # noqa """Wrap a function or context manager to support tail calls.""" __wrapped__ = None def __new__(cls, thing, *args, **kwargs): """Return a tail-call proxy.""" if isinstance(thing, cls): return thing return functools.wraps(thing)(super().__new__(cls, *args, **kwargs)) def __get__(self, instance, owner): """Return tail-call proxy if wrapped func is descriptor, else self.""" getter = inspect.getattr_static(self.__wrapped__, "__get__", None) if getter is None: return self return tail_call(getter(self.__wrapped__, instance, owner)) def __call__(self, *args, **kwargs): """Return the result of calling the wrapped function. Tail calls of functions wrapped with tail_call will be eliminated. Non-tail calls of such functions should be handled with the "call" method. """ frame = sys._getframe() # pylint: disable=protected-access reified_call = _ReifiedCall(self.__wrapped__, args, kwargs) grandparent_code = None try: grandparent_code = frame.f_back.f_back.f_code except AttributeError: pass if grandparent_code == frame.f_code: raise _Tail(reified_call) with _DeferExits() as exits: while True: try: return reified_call.function( *reified_call.args, **reified_call.kwargs ) except _Tail as tail: reified_call = tail.reified_call exits.append(tail.exits) @property def call(self): """Return the result of calling the wrapped function directly. This avoids the tail-call machinery. """ return self.__wrapped__ def __enter__(self): """Enter the wrapped context.""" return self.__wrapped__.__enter__() def __exit__(self, exc_type, exc_value, traceback): """Exit the wrapped context, unless this is a tail call. In that case, defer exiting the context until the tail calls complete. """ exit_func = self.__wrapped__.__exit__ if exc_type is not _Tail: return exit_func(exc_type, exc_value, traceback) exc_value.exits.append(exit_func) return None
Comments