hat.util

  1"""Common utility functions"""
  2
  3import contextlib
  4import inspect
  5import socket
  6import typing
  7
  8
  9T = typing.TypeVar('T')
 10
 11
 12def register_type_alias(name: str):
 13    """Register type alias
 14
 15    This function is temporary hack replacement for typing.TypeAlias.
 16
 17    It is expected that calling location will have `name` in local namespace
 18    with type value. This function will wrap that type inside `typing.TypeVar`
 19    and update annotations.
 20
 21    """
 22    frame = inspect.stack()[1][0]
 23    f_locals = frame.f_locals
 24    t = f_locals[name]
 25    f_locals[name] = typing.TypeVar(name, t, t)
 26    f_locals.setdefault('__annotations__', {})[name] = typing.Type[t]
 27
 28
 29def first(xs: typing.Iterable[T],
 30          fn: typing.Callable[[T], typing.Any] = lambda _: True,
 31          default: typing.Optional[T] = None
 32          ) -> typing.Optional[T]:
 33    """Return the first element from iterable that satisfies predicate `fn`,
 34    or `default` if no such element exists.
 35
 36    Result of predicate `fn` can be of any type. Predicate is satisfied if it's
 37    return value is truthy.
 38
 39    Args:
 40        xs: collection
 41        fn: predicate
 42        default: default value
 43
 44    Example::
 45
 46        assert first(range(3)) == 0
 47        assert first(range(3), lambda x: x > 1) == 2
 48        assert first(range(3), lambda x: x > 2) is None
 49        assert first(range(3), lambda x: x > 2, 123) == 123
 50        assert first({1: 'a', 2: 'b', 3: 'c'}) == 1
 51        assert first([], default=123) == 123
 52
 53    """
 54    return next((i for i in xs if fn(i)), default)
 55
 56
 57class RegisterCallbackHandle(typing.NamedTuple):
 58    """Handle for canceling callback registration."""
 59
 60    cancel: typing.Callable[[], None]
 61    """cancel callback registration"""
 62
 63    def __enter__(self):
 64        return self
 65
 66    def __exit__(self, *args):
 67        self.cancel()
 68
 69
 70ExceptionCb = typing.Callable[[Exception], None]
 71"""Exception callback"""
 72
 73
 74class CallbackRegistry:
 75    """Registry that enables callback registration and notification.
 76
 77    Callbacks in the registry are notified sequentially with
 78    `CallbackRegistry.notify`. If a callback raises an exception, the
 79    exception is caught and `exception_cb` handler is called. Notification of
 80    subsequent callbacks is not interrupted. If handler is `None`, the
 81    exception is reraised and no subsequent callback is notified.
 82
 83    Example::
 84
 85        x = []
 86        y = []
 87        registry = CallbackRegistry()
 88
 89        registry.register(x.append)
 90        registry.notify(1)
 91
 92        with registry.register(y.append):
 93            registry.notify(2)
 94
 95        registry.notify(3)
 96
 97        assert x == [1, 2, 3]
 98        assert y == [2]
 99
100    """
101
102    def __init__(self,
103                 exception_cb: typing.Optional[ExceptionCb] = None):
104        self._exception_cb = exception_cb
105        self._cbs = []  # type: typing.List[typing.Callable]
106
107    def register(self,
108                 cb: typing.Callable
109                 ) -> RegisterCallbackHandle:
110        """Register a callback."""
111        self._cbs.append(cb)
112        return RegisterCallbackHandle(lambda: self._cbs.remove(cb))
113
114    def notify(self, *args, **kwargs):
115        """Notify all registered callbacks."""
116        for cb in self._cbs:
117            try:
118                cb(*args, **kwargs)
119            except Exception as e:
120                if self._exception_cb:
121                    self._exception_cb(e)
122                else:
123                    raise
124
125
126def get_unused_tcp_port(host: str = '127.0.0.1') -> int:
127    """Search for unused TCP port"""
128    with contextlib.closing(socket.socket()) as sock:
129        sock.bind((host, 0))
130        return sock.getsockname()[1]
131
132
133def get_unused_udp_port(host: str = '127.0.0.1') -> int:
134    """Search for unused UDP port"""
135    with contextlib.closing(socket.socket(type=socket.SOCK_DGRAM)) as sock:
136        sock.bind((host, 0))
137        return sock.getsockname()[1]
138
139
140# HACK type alias
141register_type_alias('ExceptionCb')