omlish 0.0.0.dev4__tar.gz → 0.0.0.dev6__tar.gz

{omlish-0.0.0.dev4/omlish.egg-info → omlish-0.0.0.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omlish
-Version: 0.0.0.dev4
+Version: 0.0.0.dev6
 Summary: omlish
 Author: wrmsr
 License: BSD-3-Clause

{omlish-0.0.0.dev4 → omlish-0.0.0.dev6}/omlish/__about__.py

@@ -3,4 +3,4 @@ __url__ = 'https://github.com/wrmsr/omlish'
 __license__ = 'BSD-3-Clause'
 __requires_python__ = '>=3.12'
 
-__version__ = '0.0.0.dev4'
+__version__ = '0.0.0.dev6'

{omlish-0.0.0.dev4 → omlish-0.0.0.dev6}/omlish/__init__.py

@@ -5,4 +5,4 @@ REQUIRED_PYTHON_VERSION = (3, 12)
 
 
 if _sys.version_info < REQUIRED_PYTHON_VERSION:
-    raise EnvironmentError(f'Python version {_sys.version_info=} < {REQUIRED_PYTHON_VERSION}')
+    raise RuntimeError(f'Python version {_sys.version_info=} < {REQUIRED_PYTHON_VERSION}')

{omlish-0.0.0.dev4 → omlish-0.0.0.dev6}/omlish/asyncs/__init__.py

@@ -6,6 +6,15 @@ from .asyncs import (  # noqa
     syncable_iterable,
 )
 
+from .bridge import (  # noqa
+    a_to_s,
+    is_in_bridge,
+    s_to_a,
+    s_to_a_await,
+    trivial_a_to_s,
+    trivial_s_to_a,
+)
+
 from .flavors import (  # noqa
     ContextManagerAdapter,
     Flavor,
@@ -22,6 +31,7 @@ from .flavors import (  # noqa
     mark_asyncio,
     mark_flavor,
     mark_trio,
+    with_adapter_loop,
 )
 
 from .futures import (  # noqa
@@ -32,7 +42,3 @@ from .futures import (  # noqa
     wait_dependent_futures,
     wait_futures,
 )
-
-from .trio_asyncio import (  # noqa
-    with_trio_asyncio_loop,
-)

omlish-0.0.0.dev6/omlish/asyncs/anyio.py

@@ -0,0 +1,216 @@
+"""
+TODO:
+ - bane
+  - owned lock
+  - async once
+
+lookit:
+ - https://github.com/davidbrochart/sqlite-anyio/blob/a3ba4c6ef0535b14a5a60071fcd6ed565a514963/sqlite_anyio/sqlite.py
+ - https://github.com/rafalkrupinski/ratelimit-anyio/blob/2910a8a3d6fa54ed17ee6ba457686c9f7a4c4beb/src/ratelimit_anyio/__init__.py
+ - https://github.com/nekitdev/async-extensions/tree/main/async_extensions
+ - https://github.com/kinnay/anynet/tree/master/anynet
+ - https://github.com/M-o-a-T/asyncscope
+ - https://github.com/M-o-a-T/aevent
+ - https://github.com/florimondmanca/aiometer
+ - https://github.com/sanitizers/octomachinery/blob/b36c3d3d49da813ac46e361424132955a4e99ac8/octomachinery/utils/asynctools.py
+"""  # noqa
+import typing as ta
+
+import anyio.streams.memory
+import anyio.streams.stapled
+import sniffio
+
+from .. import check
+from .. import lang
+
+
+T = ta.TypeVar('T')
+
+
+##
+
+
+async def anyio_eof_to_empty(fn: ta.Callable[..., ta.Awaitable[T]], *args: ta.Any, **kwargs: ta.Any) -> T | bytes:
+    try:
+        return await fn(*args, **kwargs)
+    except anyio.EndOfStream:
+        return b''
+
+
+async def gather(*fns: ta.Callable[..., ta.Awaitable[T]], take_first: bool = False) -> list[lang.Maybe[T]]:
+    results: list[lang.Maybe[T]] = [lang.empty()] * len(fns)
+
+    async def inner(fn, i):
+        results[i] = lang.just(await fn())
+        if take_first:
+            tg.cancel_scope.cancel()
+
+    async with anyio.create_task_group() as tg:
+        for i, fn in enumerate(fns):
+            tg.start_soon(inner, fn, i)
+
+    return results
+
+
+async def first(*fns: ta.Callable[..., ta.Awaitable[T]], **kwargs: ta.Any) -> list[lang.Maybe[T]]:
+    return await gather(*fns, take_first=True, **kwargs)
+
+
+##
+
+
+def get_current_task() -> anyio.TaskInfo | None:
+    try:
+        return anyio.get_current_task()
+    except sniffio.AsyncLibraryNotFoundError:
+        return None
+
+
+#
+
+
+BackendTask: ta.TypeAlias = ta.Union[  # noqa
+    # asyncio.tasks.Task,
+    # trio.lowlevel.Task,
+    ta.Any,
+]
+
+
+def _is_class_named(obj: ta.Any, m: str, n: str) -> bool:
+    cls = obj.__class__
+    return cls.__module__ == m and cls.__name__ == n
+
+
+def get_backend_task(at: anyio.TaskInfo) -> BackendTask | None:
+    if _is_class_named(at, 'anyio._backends._asyncio', 'AsyncIOTaskInfo'):
+        # https://github.com/agronholm/anyio/blob/8907964926a24461840eee0925d3f355e729f15d/src/anyio/_backends/_asyncio.py#L1846  # noqa
+        # weakref.ref
+        obj = at._task()  # type: ignore  # noqa
+        if obj is not None and not (
+                _is_class_named(obj, '_asyncio', 'Task') or
+                _is_class_named(obj, 'asyncio.tasks', 'Task')
+        ):
+            raise TypeError(obj)
+        return obj
+
+    elif _is_class_named(at, 'anyio._backends._trio', 'TrioTaskInfo'):
+        # https://github.com/agronholm/anyio/blob/8907964926a24461840eee0925d3f355e729f15d/src/anyio/_backends/_trio.py#L850  # noqa
+        # weakref.proxy
+        # https://stackoverflow.com/a/62144308 :|
+        obj = at._task.__repr__.__self__  # type: ignore  # noqa
+        if obj is not None and not _is_class_named(obj, 'trio.lowlevel', 'Task'):
+            raise TypeError(obj)
+        return obj
+
+    else:
+        raise TypeError(at)
+
+
+def get_current_backend_task() -> BackendTask | None:
+    if (at := get_current_task()) is not None:
+        return get_backend_task(at)
+    else:
+        return None
+
+
+##
+
+
+def split_memory_object_streams(
+        *args: anyio.create_memory_object_stream[T],
+) -> tuple[
+    anyio.streams.memory.MemoryObjectSendStream[T],
+    anyio.streams.memory.MemoryObjectReceiveStream[T],
+]:
+    [tup] = args
+    return tup
+
+
+# FIXME: https://github.com/python/mypy/issues/15238
+# FIXME: https://youtrack.jetbrains.com/issues?q=tag:%20%7BPEP%20695%7D
+def create_memory_object_stream[T](max_buffer_size: float = 0) -> tuple[
+    anyio.streams.memory.MemoryObjectSendStream[T],
+    anyio.streams.memory.MemoryObjectReceiveStream[T],
+]:
+    return anyio.create_memory_object_stream[T](max_buffer_size)  # noqa
+
+
+def staple_memory_object_stream(
+        *args: anyio.create_memory_object_stream[T],
+) -> anyio.streams.stapled.StapledObjectStream[T]:
+    send, receive = args
+    return anyio.streams.stapled.StapledObjectStream(
+        check.isinstance(send, anyio.streams.memory.MemoryObjectSendStream),  # type: ignore
+        check.isinstance(receive, anyio.streams.memory.MemoryObjectReceiveStream),  # type: ignore
+    )
+
+
+# FIXME: https://github.com/python/mypy/issues/15238
+# FIXME: https://youtrack.jetbrains.com/issues?q=tag:%20%7BPEP%20695%7D
+def staple_memory_object_stream2[T](max_buffer_size: float = 0) -> anyio.streams.stapled.StapledObjectStream[T]:
+    send, receive = anyio.create_memory_object_stream[T](max_buffer_size)
+    return anyio.streams.stapled.StapledObjectStream(
+        check.isinstance(send, anyio.streams.memory.MemoryObjectSendStream),  # type: ignore
+        check.isinstance(receive, anyio.streams.memory.MemoryObjectReceiveStream),  # type: ignore
+    )
+
+
+##
+
+
+class Once:
+    def __init__(self) -> None:
+        super().__init__()
+        self._done = False
+        self._lock = anyio.Lock()
+
+    async def do(self, fn: ta.Callable[[], ta.Awaitable[None]]) -> bool:
+        if self._done:
+            return False
+        async with self._lock:
+            if self._done:
+                return False  # type: ignore
+            try:
+                await fn()
+            finally:
+                self._done = True
+            return True
+
+
+class Lazy(ta.Generic[T]):
+    def __init__(self) -> None:
+        super().__init__()
+        self._once = Once()
+        self._v: lang.Maybe[T] = lang.empty()
+
+    def peek(self) -> lang.Maybe[T]:
+        return self._v
+
+    def set(self, v: T) -> None:
+        self._v = lang.just(v)
+
+    async def get(self, fn: ta.Callable[[], ta.Awaitable[T]]) -> T:
+        async def do():
+            self._v = lang.just(await fn())
+        await self._once.do(do)
+        return self._v.must()
+
+
+class LazyFn(ta.Generic[T]):
+    def __init__(self, fn: ta.Callable[[], ta.Awaitable[T]]) -> None:
+        super().__init__()
+        self._fn = fn
+        self._once = Once()
+        self._v: lang.Maybe[T] = lang.empty()
+
+    def peek(self) -> lang.Maybe[T]:
+        return self._v
+
+    def set(self, v: T) -> None:
+        self._v = lang.just(v)
+
+    async def get(self) -> T:
+        async def do():
+            self._v = lang.just(await self._fn())
+        await self._once.do(do)
+        return self._v.must()

omlish-0.0.0.dev6/omlish/asyncs/asyncio.py

@@ -0,0 +1,42 @@
+import asyncio
+import contextlib
+import functools
+import typing as ta
+
+
+CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
+
+
+def asyncio_once(fn: CallableT) -> CallableT:
+    future = None
+
+    @functools.wraps(fn)
+    async def inner(*args, **kwargs):
+        nonlocal future
+        if not future:
+            future = asyncio.create_task(fn(*args, **kwargs))
+        return await future
+
+    return ta.cast(CallableT, inner)
+
+
+def get_real_current_loop() -> asyncio.AbstractEventLoop | None:
+    return asyncio.get_event_loop_policy()._local._loop  # type: ignore  # noqa
+
+
+def drain_tasks(loop=None):
+    if loop is None:
+        loop = get_real_current_loop()
+
+    while loop._ready or loop._scheduled:  # noqa
+        loop._run_once()  # noqa
+
+
+@contextlib.contextmanager
+def draining_asyncio_tasks() -> ta.Iterator[None]:
+    loop = get_real_current_loop()
+    try:
+        yield
+    finally:
+        if loop is not None:
+            drain_tasks(loop)  # noqa

{omlish-0.0.0.dev4 → omlish-0.0.0.dev6}/omlish/asyncs/asyncs.py

@@ -5,10 +5,10 @@ TODO:
   517 ns ± 13.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
   - injected io provider - sync vs greenlet aio trampolined
  - push/pull bridge?
+ - move to lang
 
 https://github.com/sqlalchemy/sqlalchemy/blob/1e75c189da721395bc8c2d899c722a5b9a170404/lib/sqlalchemy/util/_concurrency_py3k.py#L83
 """
-import contextlib
 import functools
 import typing as ta
 
@@ -16,8 +16,7 @@ import typing as ta
 T = ta.TypeVar('T')
 
 
-def sync_await(fn: ta.Callable[..., T], *args, **kwargs) -> T:
-    ret: ta.Any
+def sync_await(fn: ta.Callable[..., T], *args: ta.Any, **kwargs: ta.Any) -> T:
     ret = missing = object()
 
     async def gate():
@@ -25,13 +24,17 @@ def sync_await(fn: ta.Callable[..., T], *args, **kwargs) -> T:
         ret = await fn(*args, **kwargs)  # type: ignore
 
     cr = gate()
-    with contextlib.closing(cr):
-        with contextlib.suppress(StopIteration):
+    try:
+        try:
             cr.send(None)
+        except StopIteration:
+            pass
         if ret is missing or cr.cr_await is not None or cr.cr_running:
             raise TypeError('Not terminated')
+    finally:
+        cr.close()
 
-    return ta.cast(T, ret)
+    return ret  # type: ignore
 
 
 def sync_list(fn: ta.Callable[..., ta.AsyncIterator[T]], *args, **kwargs) -> list[T]:

omlish-0.0.0.dev6/omlish/asyncs/bridge.py

@@ -0,0 +1,316 @@
+"""
+https://github.com/kubernetes/kubernetes/blob/60c4c2b2521fb454ce69dee737e3eb91a25e0535/pkg/controller/volume/persistentvolume/pv_controller.go#L60-L63
+
+==================================================================
+PLEASE DO NOT ATTEMPT TO SIMPLIFY THIS CODE.
+KEEP THE SPACE SHUTTLE FLYING.
+==================================================================
+
+TODO:
+ - reuse greenlet if nested somehow?
+"""
+import itertools
+import sys
+import types
+import typing as ta
+import weakref
+
+from .. import lang
+from .asyncs import sync_await
+
+
+if ta.TYPE_CHECKING:
+    import asyncio
+
+    import greenlet
+
+    from . import anyio as aiu
+
+else:
+    asyncio = lang.proxy_import('asyncio')
+
+    greenlet = lang.proxy_import('greenlet')
+
+    aiu = lang.proxy_import('.anyio', __package__)
+
+
+T = ta.TypeVar('T')
+
+
+##
+
+
+def trivial_s_to_a(fn):
+    async def inner(*args, **kwargs):
+        return fn(*args, **kwargs)
+    return inner
+
+
+def trivial_a_to_s(fn):
+    def inner(*args, **kwargs):
+        return sync_await(fn, *args, **kwargs)
+    return inner
+
+
+##
+# https://gist.github.com/snaury/202bf4f22c41ca34e56297bae5f33fef
+
+
+class BridgeAwaitRequiredError(Exception):
+    pass
+
+
+class MissingBridgeGreenletError(Exception):
+    pass
+
+
+class UnexpectedBridgeNestingError(Exception):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # breakpoint()
+
+
+#
+
+
+_DEBUG_PRINT: ta.Callable[..., None] | None = None
+# _DEBUG_PRINT = print  # noqa
+
+_TRACK_TRANSITION_OBJS = False
+
+
+#
+
+
+_BRIDGE_TRANSITIONS_SEQ = itertools.count()
+
+
+class _BridgeTransition(ta.NamedTuple):
+    seq: int
+    a_to_s: bool
+
+    obj_cls: type
+    obj_id: int
+
+    obj: ta.Any
+
+
+def _make_transition(seq: int, a_to_s: bool, obj: ta.Any) -> _BridgeTransition:
+    return _BridgeTransition(seq, a_to_s, obj.__class__, id(obj), (obj if _TRACK_TRANSITION_OBJS else None))
+
+
+_BRIDGED_TASKS: ta.MutableMapping[ta.Any, list[_BridgeTransition]] = weakref.WeakKeyDictionary()
+
+_BRIDGE_GREENLET_ATTR = f'__{__package__.replace(".", "__")}__bridge_greenlet__'
+
+
+def _push_transition(a_to_s: bool, l: list[_BridgeTransition], t: _BridgeTransition) -> _BridgeTransition:
+    l.append(t)
+    if _DEBUG_PRINT:
+        _DEBUG_PRINT(f'_push_transition: {a_to_s=} {id(l)=} {t=}')
+    return t
+
+
+def _pop_transition(a_to_s: bool, l: list[_BridgeTransition]) -> _BridgeTransition:
+    t = l.pop()
+    if _DEBUG_PRINT:
+        _DEBUG_PRINT(f'_pop_transition: {a_to_s=} {id(l)=} {t=}')
+    return t
+
+
+def _get_transitions() -> list[_BridgeTransition]:
+    l: list[_BridgeTransition] = []
+
+    if (t := aiu.get_current_backend_task()) is not None:
+        try:
+            tl = _BRIDGED_TASKS[t]
+        except KeyError:
+            pass
+        else:
+            l.extend(tl)
+
+    g = greenlet.getcurrent()
+    try:
+        gl = getattr(g, _BRIDGE_GREENLET_ATTR)
+    except AttributeError:
+        pass
+    else:
+        l.extend(gl)
+
+    l.sort(key=lambda t: (t.seq, t.a_to_s))
+    return l
+
+
+def is_in_bridge() -> bool:
+    if _DEBUG_PRINT:
+        _DEBUG_PRINT(_get_transitions())
+
+    if (t := aiu.get_current_backend_task()) is not None:
+        try:
+            tl = _BRIDGED_TASKS[t]
+        except KeyError:
+            last_t = None
+        else:
+            if tl:
+                last_t = tl[-1]
+            else:
+                last_t = None
+    else:
+        last_t = None
+
+    g = greenlet.getcurrent()
+    try:
+        gl = getattr(g, _BRIDGE_GREENLET_ATTR)
+    except AttributeError:
+        last_g = None
+    else:
+        if gl:
+            last_g = gl[-1]
+        else:
+            last_g = None
+
+    if last_t is None:
+        if last_g is None:
+            return False
+        o = last_g
+    else:  # noqa
+        if last_g is None or last_g.seq < last_t.seq:
+            o = last_t
+        else:
+            o = last_g
+
+    in_a = (t is not None)
+
+    if _DEBUG_PRINT:
+        _DEBUG_PRINT(f'{o.a_to_s=} {in_a=}')
+
+    return in_a != o.a_to_s
+
+
+def _safe_cancel_awaitable(awaitable: ta.Awaitable[ta.Any]) -> None:
+    # https://docs.python.org/3/reference/datamodel.html#coroutine.close
+    if asyncio.iscoroutine(awaitable):
+        awaitable.close()  # noqa
+
+
+def s_to_a_await(awaitable: ta.Awaitable[T]) -> T:
+    g = greenlet.getcurrent()
+
+    if not getattr(g, _BRIDGE_GREENLET_ATTR, False):
+        _safe_cancel_awaitable(awaitable)
+        raise MissingBridgeGreenletError
+
+    return g.parent.switch(awaitable)
+
+
+def s_to_a(fn, *, require_await=False):
+    @types.coroutine
+    def outer(*args, **kwargs):
+        def inner():
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                if (gl2 := getattr(g, _BRIDGE_GREENLET_ATTR)) is not gl:  # noqa
+                    raise UnexpectedBridgeNestingError
+                if (cur_g := _pop_transition(False, gl)) is not added_g:  # noqa
+                    raise UnexpectedBridgeNestingError
+                if gl:
+                    raise UnexpectedBridgeNestingError
+
+        seq = next(_BRIDGE_TRANSITIONS_SEQ)
+
+        g = greenlet.greenlet(inner)
+        setattr(g, _BRIDGE_GREENLET_ATTR, gl := [])  # type: ignore
+        added_g = _push_transition(False, gl, _make_transition(seq, False, g))
+
+        if (t := aiu.get_current_backend_task()) is not None:
+            try:
+                tl = _BRIDGED_TASKS[t]
+            except KeyError:
+                tl = _BRIDGED_TASKS[t] = []
+            added_t = _push_transition(False, tl, _make_transition(seq, False, g))
+
+        try:
+            result: ta.Any = g.switch()
+            switch_occurred = False
+            while not g.dead:
+                switch_occurred = True
+                try:
+                    value = yield result
+                except BaseException:  # noqa
+                    result = g.throw(*sys.exc_info())
+                else:
+                    result = g.switch(value)
+
+            if require_await and not switch_occurred:
+                raise BridgeAwaitRequiredError
+
+            return result
+
+        finally:
+            if t is not None:
+                if (tl2 := _BRIDGED_TASKS[t]) is not tl:  # noqa
+                    raise UnexpectedBridgeNestingError
+                if (cur_t := _pop_transition(False, tl)) is not added_t:  # noqa
+                    raise UnexpectedBridgeNestingError
+
+    return outer
+
+
+def a_to_s(fn):
+    def inner(*args, **kwargs):
+        seq = next(_BRIDGE_TRANSITIONS_SEQ)
+
+        if (t := aiu.get_current_backend_task()) is not None:
+            try:
+                tl = _BRIDGED_TASKS[t]
+            except KeyError:
+                tl = _BRIDGED_TASKS[t] = []
+            added_t = _push_transition(True, tl, _make_transition(seq, True, t))
+        else:
+            added_t = None
+
+        g = greenlet.getcurrent()
+        try:
+            gl = getattr(g, _BRIDGE_GREENLET_ATTR)
+        except AttributeError:
+            setattr(g, _BRIDGE_GREENLET_ATTR, gl := [])
+        added_g = _push_transition(True, gl, _make_transition(seq, True, g))
+
+        try:
+            ret = missing = object()
+
+            async def gate():
+                nonlocal ret
+                ret = await fn(*args, **kwargs)
+
+            cr = gate()
+            sv = None
+            try:
+                while True:
+                    try:
+                        sv = cr.send(sv)
+                    except StopIteration:
+                        break
+
+                    if ret is missing or cr.cr_await is not None or cr.cr_running:
+                        sv = s_to_a_await(sv)
+
+            finally:
+                cr.close()
+
+        finally:
+            if t is not None:
+                if (tl2 := _BRIDGED_TASKS[t]) is not tl:  # noqa
+                    raise UnexpectedBridgeNestingError
+                if (cur_t := _pop_transition(True, tl)) is not added_t:  # noqa
+                    raise UnexpectedBridgeNestingError
+
+            if (gl2 := getattr(g, _BRIDGE_GREENLET_ATTR)) is not gl:  # noqa
+                raise UnexpectedBridgeNestingError
+            if (cur_g := _pop_transition(True, gl)) is not added_g:  # noqa
+                raise UnexpectedBridgeNestingError
+
+        return ret
+
+    return inner