codekingpro/portable-devtools

1"""2Utility decorators that help build state machines3"""4 5import functools6 7from mitmproxy.proxy import events8 9 10def expect(*event_types):11    """12    Only allow the given event type.13    If another event is passed, an AssertionError is raised.14    """15 16    def decorator(f):17        if __debug__ is True:18 19            @functools.wraps(f)20            def _check_event_type(self, event: events.Event):21                if isinstance(event, event_types):22                    return f(self, event)23                else:24                    event_types_str = (25                        "|".join(e.__name__ for e in event_types) or "no events"26                    )27                    raise AssertionError(28                        f"Unexpected event type at {f.__qualname__}: "29                        f"Expected {event_types_str}, got {event}."30                    )31 32            return _check_event_type33        else:  # pragma: no cover34            return f35 36    return decorator37 38 39class ReceiveBuffer:40    """41    A data structure to collect stream contents efficiently in O(n).42    """43 44    _chunks: list[bytes]45    _len: int46 47    def __init__(self):48        self._chunks = []49        self._len = 050 51    def __iadd__(self, other: bytes):52        assert isinstance(other, bytes)53        self._chunks.append(other)54        self._len += len(other)55        return self56 57    def __len__(self):58        return self._len59 60    def __bytes__(self):61        return b"".join(self._chunks)62 63    def __bool__(self):64        return self._len > 065 66    def clear(self):67        self._chunks.clear()68        self._len = 069