omlish 0.0.0.dev5__py3-none-any.whl → 0.0.0.dev6__py3-none-any.whl

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.dev5'
+__version__ = '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,

omlish/asyncs/anyio.py

@@ -18,6 +18,7 @@ import typing as ta
 
 import anyio.streams.memory
 import anyio.streams.stapled
+import sniffio
 
 from .. import check
 from .. import lang
@@ -26,6 +27,9 @@ 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)
@@ -33,6 +37,85 @@ async def anyio_eof_to_empty(fn: ta.Callable[..., ta.Awaitable[T]], *args: ta.An
         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[
@@ -72,25 +155,6 @@ def staple_memory_object_stream2[T](max_buffer_size: float = 0) -> anyio.streams
     )
 
 
-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)
-
-
 ##
 
 

omlish/asyncs/asyncio.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import functools
 import typing as ta
 
@@ -17,3 +18,25 @@ def asyncio_once(fn: CallableT) -> CallableT:
         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/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/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

omlish/asyncs/trio_asyncio.py

@@ -9,6 +9,7 @@ if ta.TYPE_CHECKING:
 
     import sniffio
     import trio_asyncio
+
 else:
     asyncio = lang.proxy_import('asyncio')
 
@@ -21,7 +22,7 @@ def check_trio_asyncio() -> None:
         raise RuntimeError('trio_asyncio loop not running')
 
 
-def with_trio_asyncio_loop(*, wait=False):
+def with_trio_asyncio_loop(*, wait=False, strict=False):
     def outer(fn):
         @functools.wraps(fn)
         async def inner(*args, **kwargs):
@@ -30,7 +31,10 @@ def with_trio_asyncio_loop(*, wait=False):
                 return
 
             if sniffio.current_async_library() != 'trio':
-                raise RuntimeError('trio loop not running')
+                if strict:
+                    raise RuntimeError('trio loop not running')
+                await fn(*args, **kwargs)
+                return
 
             loop: asyncio.BaseEventLoop
             async with trio_asyncio.open_loop() as loop:
@@ -40,7 +44,7 @@ def with_trio_asyncio_loop(*, wait=False):
                     if wait:
                         # FIXME: lol
                         while asyncio.all_tasks(loop):
-                            await asyncio.sleep(.2)
+                            await asyncio.sleep(.1)
 
         return inner
 

omlish/collections/__init__.py

@@ -37,6 +37,7 @@ from .identity import (  # noqa
     IdentityKeyDict,
     IdentitySet,
     IdentityWrapper,
+    IdentityWeakSet,
 )
 
 from .indexed import (  # noqa

omlish/collections/identity.py

@@ -1,6 +1,7 @@
 import contextlib
 import operator as op
 import typing as ta
+import weakref
 
 from .. import lang
 from .mappings import yield_dict_init
@@ -103,3 +104,9 @@ class IdentitySet(ta.MutableSet[T]):
 
     def __iter__(self) -> ta.Iterator[T]:
         return iter(self._dict.values())
+
+
+class IdentityWeakSet(weakref.WeakSet):
+    def __init__(self, init=None):
+        super().__init__()
+        self.data = IdentitySet(init)  # type: ignore

omlish/configs/strings.py

@@ -0,0 +1,94 @@
+"""
+TODO:
+ - reflecty generalized rewriter, obviously..
+"""
+import collections.abc
+import typing as ta
+
+from .. import dataclasses as dc
+from .. import lang
+from ..text import glyphsplit
+
+
+T = ta.TypeVar('T')
+
+
+class InterpolateStringsMetadata(lang.Marker):
+    pass
+
+
+@dc.field_modifier
+def secret_or_key_field(f: dc.Field) -> dc.Field:
+    return dc.update_field_metadata(f, {
+        InterpolateStringsMetadata: True,
+    })
+
+
+@dc.field_modifier
+def interpolate_field(f: dc.Field) -> dc.Field:
+    return dc.update_field_metadata(f, {InterpolateStringsMetadata: True})
+
+
+class StringRewriter:
+    def __init__(self, fn: ta.Callable[[str], str]) -> None:
+        super().__init__()
+        self._fn = fn
+
+    def __call__(self, v: T, *, _soft: bool = False) -> T:
+        if v is None:
+            return None  # type: ignore
+
+        if dc.is_dataclass(v):
+            kw = {}
+            for f in dc.fields(v):
+                fv = getattr(v, f.name)
+                nfv = self(fv, _soft=not f.metadata.get(InterpolateStringsMetadata))
+                if fv is not nfv:
+                    kw[f.name] = nfv
+            if not kw:
+                return v  # type: ignore
+            return dc.replace(v, **kw)
+
+        if isinstance(v, str):
+            if not _soft:
+                v = self._fn(v)  # type: ignore
+            return v  # type: ignore
+
+        if isinstance(v, lang.BUILTIN_SCALAR_ITERABLE_TYPES):
+            return v  # type: ignore
+
+        if isinstance(v, collections.abc.Mapping):
+            nm = []
+            b = False
+            for mk, mv in v.items():
+                nk, nv = self(mk, _soft=_soft), self(mv, _soft=_soft)
+                nm.append((nk, nv))
+                b |= nk is not mk or nv is not mv
+            if not b:
+                return v  # type: ignore
+            return v.__class__(nm)  # type: ignore
+
+        if isinstance(v, (collections.abc.Sequence, collections.abc.Set)):
+            nl = []
+            b = False
+            for le in v:
+                ne = self(le, _soft=_soft)
+                nl.append(ne)
+                b |= ne is not le
+            if not b:
+                return v  # type: ignore
+            return v.__class__(nl)  # type: ignore
+
+        return v
+
+
+def interpolate_strings(v: T, rpl: ta.Mapping[str, str]) -> T:
+    def fn(v):
+        if not v:
+            return v
+        sps = glyphsplit.split_braces(v)
+        if len(sps) == 1 and isinstance(sps[0], str):
+            return sps[0]
+        return ''.join(rpl[p.s] if isinstance(p, glyphsplit.GlyphMatch) else p for p in sps)
+
+    return StringRewriter(fn)(v)