mwchase icon

Extremely normal Python code

mwchase | PRO | 02/27/19 12:34:23 AM UTC | 0 ⭐ | 830 👁️ | Never ⏰ | []
Python |

3.42 KB

|

None

|

0 👍

/

0 👎

"""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