ommlds 0.0.0.dev440__py3-none-any.whl → 0.0.0.dev480__py3-none-any.whl

ommlds/minichain/resources.py

@@ -3,7 +3,6 @@ TODO:
  - unify with omlish.sql.api.resources -> omlish.resources
 """
 import contextlib
-import threading
 import typing as ta
 
 from omlish import check
@@ -32,6 +31,7 @@ class ResourcesRefNotRegisteredError(Exception):
     pass
 
 
+@ta.final
 class Resources(lang.Final, lang.NotPicklable):
     def __init__(
             self,
@@ -43,26 +43,25 @@ class Resources(lang.Final, lang.NotPicklable):
 
         self._no_autoclose = no_autoclose
 
-        self._lock = threading.RLock()
-
         self._closed = False
 
         self._refs: ta.MutableSet[ResourcesRef] = col.IdentitySet()
 
-        self._es = contextlib.ExitStack()
-        self._es.__enter__()
+        self._aes = contextlib.AsyncExitStack()
 
         if init_ref is not None:
             self.add_ref(init_ref)
 
+    async def init(self) -> None:
+        await self._aes.__aenter__()
+
     @property
     def autoclose(self) -> bool:
         return not self._no_autoclose
 
     @property
     def num_refs(self) -> int:
-        with self._lock:
-            return len(self._refs)
+        return len(self._refs)
 
     @property
     def closed(self) -> bool:
@@ -77,15 +76,16 @@ class Resources(lang.Final, lang.NotPicklable):
         pass
 
     @classmethod
-    def new(cls, **kwargs: ta.Any) -> ta.ContextManager['Resources']:
-        @contextlib.contextmanager
-        def inner():
+    def new(cls, **kwargs: ta.Any) -> ta.AsyncContextManager['Resources']:
+        @contextlib.asynccontextmanager
+        async def inner():
             init_ref = Resources._InitRef()
             res = Resources(init_ref=init_ref, **kwargs)
+            await res.init()
             try:
                 yield res
             finally:
-                res.remove_ref(init_ref)
+                await res.remove_ref(init_ref)
 
         return inner()
 
@@ -93,30 +93,30 @@ class Resources(lang.Final, lang.NotPicklable):
 
     def add_ref(self, ref: ResourcesRef) -> None:
         check.isinstance(ref, ResourcesRef)
-        with self._lock:
-            check.state(not self._closed)
-            self._refs.add(ref)
+        check.state(not self._closed)
+        self._refs.add(ref)
 
     def has_ref(self, ref: ResourcesRef) -> bool:
-        with self._lock:
-            return ref in self._refs
+        return ref in self._refs
 
-    def remove_ref(self, ref: ResourcesRef) -> None:
+    async def remove_ref(self, ref: ResourcesRef) -> None:
         check.isinstance(ref, ResourcesRef)
-        with self._lock:
-            try:
-                self._refs.remove(ref)
-            except KeyError:
-                raise ResourcesRefNotRegisteredError(ref) from None
-            if not self._no_autoclose and not self._refs:
-                self.close()
+        try:
+            self._refs.remove(ref)
+        except KeyError:
+            raise ResourcesRefNotRegisteredError(ref) from None
+        if not self._no_autoclose and not self._refs:
+            await self.aclose()
 
     #
 
     def enter_context(self, cm: ta.ContextManager[T]) -> T:
-        with self._lock:
-            check.state(not self._closed)
-            return self._es.enter_context(cm)
+        check.state(not self._closed)
+        return self._aes.enter_context(cm)
+
+    async def enter_async_context(self, cm: ta.AsyncContextManager[T]) -> T:
+        check.state(not self._closed)
+        return await self._aes.enter_async_context(cm)
 
     #
 
@@ -125,12 +125,11 @@ class Resources(lang.Final, lang.NotPicklable):
 
     #
 
-    def close(self) -> None:
-        with self._lock:
-            try:
-                self._es.__exit__(None, None, None)
-            finally:
-                self._closed = True
+    async def aclose(self) -> None:
+        try:
+            await self._aes.__aexit__(None, None, None)
+        finally:
+            self._closed = True
 
     def __del__(self) -> None:
         if not self._closed:
@@ -147,24 +146,48 @@ class Resources(lang.Final, lang.NotPicklable):
 ##
 
 
+@ta.final
 class ResourceManaged(ResourcesRef, lang.Final, lang.NotPicklable, ta.Generic[T]):
+    """
+    A class to 'handoff' a ref to a `Resources`, allowing the `Resources` to temporarily survive being passed from
+    instantiation within a callee to being `__aenter__`'d in the caller.
+
+    The ref to the `Resources` is allocated in the ctor, so the contract is that an instance of this must be immediately
+    `__aenter__`'d before doing anything else with the return value of the call. Failure to do so leaks the `Resources`.
+    """
+
     def __init__(self, v: T, resources: Resources) -> None:
         super().__init__()
 
-        self._v = v
+        self.__v = v
         self.__resources = resources
 
         resources.add_ref(self)
 
+    __state: ta.Literal['new', 'entered', 'exited'] = 'new'
+
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}<{self._v!r}>'
+        return f'{self.__class__.__name__}<{self.__v!r}, {self.__state}>'
 
-    def __enter__(self) -> T:
-        return self._v
+    async def __aenter__(self) -> T:
+        check.state(self.__state == 'new')
+        self.__state = 'entered'
+        return self.__v
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.__resources.remove_ref(self)
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        check.state(self.__state == 'entered')
+        self.__state = 'exited'
+        await self.__resources.remove_ref(self)
 
+    def __del__(self) -> None:
+        if self.__state != 'exited':
+            log.error(
+                f'{__package__}.{self.__class__.__name__}.__del__: '  # noqa
+                f'%r deleted without being entered and exited! '
+                f'resources: %s',
+                repr(self),
+                repr(self.__resources),
+            )
 
 ##
 
@@ -175,11 +198,11 @@ class ResourcesOption(Option, lang.Abstract):
 
 class UseResources(tv.UniqueScalarTypedValue[Resources], ResourcesOption, lang.Final):
     @classmethod
-    @contextlib.contextmanager
-    def or_new(cls, options: ta.Sequence[Option]) -> ta.Iterator[Resources]:
+    @contextlib.asynccontextmanager
+    async def or_new(cls, options: ta.Sequence[Option]) -> ta.AsyncGenerator[Resources]:
         if (ur := tv.as_collection(options).get(UseResources)) is not None:
-            with ResourceManaged(ur.v, ur.v) as rs:
+            async with ResourceManaged(ur.v, ur.v) as rs:
                 yield rs
         else:
-            with Resources.new() as rs:
+            async with Resources.new() as rs:
                 yield rs

ommlds/minichain/search.py

@@ -19,7 +19,7 @@ class SearchHit(lang.Final):
     title: str | None
     url: str | None
     description: str | None = None
-    snippets: lang.SequenceNotStr[str] | None
+    snippets: lang.SequenceNotStr[str] | None = None
 
 
 @dc.dataclass(frozen=True, kw_only=True)

ommlds/minichain/services/_marshal.py

@@ -40,14 +40,14 @@ class _RequestResponseMarshaler(msh.Marshaler):
     def marshal(self, ctx: msh.MarshalContext, o: ta.Any) -> msh.Value:
         tv_types_set = o._typed_values_info().tv_types_set  # noqa  # FIXME
         tv_ta = tv.TypedValues[ta.Union[*tv_types_set]]  # type: ignore
-        tv_m = ctx.make(tv_ta)
+        tv_m = ctx.marshal_factory_context.make_marshaler(tv_ta)  # FIXME:
         tv_v = check.isinstance(tv_m.marshal(ctx, o._typed_values), ta.Sequence)  # noqa
 
         if self.v_m is None:
             orty: rfl.Generic = check.isinstance(rfl.type_(rfl.get_orig_class(o)), rfl.Generic)
             check.state(orty.cls in (Request, Response))
             v_rty, tv_rty = orty.args
-            v_v = ctx.make(v_rty).marshal(ctx, o.v)
+            v_v = ctx.marshal_factory_context.make_marshaler(v_rty).marshal(ctx, o.v)  # FIXME
         else:
             v_v = self.v_m.marshal(ctx, o.v)
 
@@ -57,26 +57,30 @@ class _RequestResponseMarshaler(msh.Marshaler):
         }
 
 
-class _RequestResponseMarshalerFactory(msh.SimpleMarshalerFactory):
-    def guard(self, ctx: msh.MarshalContext, rty: rfl.Type) -> bool:
-        return _is_rr_rty(rty)
+class _RequestResponseMarshalerFactory(msh.MarshalerFactory):
+    def make_marshaler(self, ctx: msh.MarshalFactoryContext, rty: rfl.Type) -> ta.Callable[[], msh.Marshaler] | None:
+        if not _is_rr_rty(rty):
+            return None
 
-    def fn(self, ctx: msh.MarshalContext, rty: rfl.Type) -> msh.Marshaler:
         if isinstance(rty, type):
             rty = rfl.type_(rfl.get_orig_class(rty))
-        if isinstance(rty, rfl.Generic):
-            v_rty, tv_rty = rty.args
-        else:
-            # FIXME: ...
-            raise TypeError(rty)
-        v_m: msh.Marshaler | None = None
-        if not isinstance(v_rty, ta.TypeVar):
-            v_m = ctx.make(v_rty)
-        return _RequestResponseMarshaler(
-            rty,
-            _get_tv_fld(rty),
-            v_m,
-        )
+
+        def inner() -> msh.Marshaler:
+            if isinstance(rty, rfl.Generic):
+                v_rty, tv_rty = rty.args
+            else:
+                # FIXME: ...
+                raise TypeError(rty)
+            v_m: msh.Marshaler | None = None
+            if not isinstance(v_rty, ta.TypeVar):
+                v_m = ctx.make_marshaler(v_rty)
+            return _RequestResponseMarshaler(
+                rty,
+                _get_tv_fld(rty),
+                v_m,
+            )
+
+        return inner
 
 
 #
@@ -108,33 +112,36 @@ class _RequestResponseUnmarshaler(msh.Unmarshaler):
         return cty(v, tvs)  # type: ignore
 
 
-class _RequestResponseUnmarshalerFactory(msh.SimpleUnmarshalerFactory):
-    def guard(self, ctx: msh.UnmarshalContext, rty: rfl.Type) -> bool:
-        return _is_rr_rty(rty)
+class _RequestResponseUnmarshalerFactory(msh.UnmarshalerFactory):
+    def make_unmarshaler(self, ctx: msh.UnmarshalFactoryContext, rty: rfl.Type) -> ta.Callable[[], msh.Unmarshaler] | None:  # noqa
+        if not _is_rr_rty(rty):
+            return None
 
-    def fn(self, ctx: msh.UnmarshalContext, rty: rfl.Type) -> msh.Unmarshaler:
-        if isinstance(rty, rfl.Generic):
-            v_rty, tv_rty = rty.args
-        else:
-            # FIXME: ...
-            raise TypeError(rty)
-        tv_types_set = check.isinstance(tv_rty, rfl.Union).args
-        tv_ta = tv.TypedValues[ta.Union[*tv_types_set]]  # type: ignore
-        tv_u = ctx.make(tv_ta)
-        v_u = ctx.make(v_rty)
-        return _RequestResponseUnmarshaler(
-            rty,
-            _get_tv_fld(rty),
-            v_u,
-            tv_u,
-        )
+        def inner() -> msh.Unmarshaler:
+            if isinstance(rty, rfl.Generic):
+                v_rty, tv_rty = rty.args
+            else:
+                # FIXME: ...
+                raise TypeError(rty)
+            tv_types_set = check.isinstance(tv_rty, rfl.Union).args
+            tv_ta = tv.TypedValues[ta.Union[*tv_types_set]]  # type: ignore
+            tv_u = ctx.make_unmarshaler(tv_ta)
+            v_u = ctx.make_unmarshaler(v_rty)
+            return _RequestResponseUnmarshaler(
+                rty,
+                _get_tv_fld(rty),
+                v_u,
+                tv_u,
+            )
+
+        return inner
 
 
 ##
 
 
 @lang.static_init
-def _install_standard_marshalling() -> None:
+def _install_standard_marshaling() -> None:
     msh.install_standard_factories(
         _RequestResponseMarshalerFactory(),
         _RequestResponseUnmarshalerFactory(),

ommlds/minichain/services/facades.py

@@ -44,15 +44,15 @@ class ServiceFacade(
         ],
     ]
 
-    def invoke(self, request: Request[RequestV, OptionT]) -> Response[ResponseV, OutputT]:
+    def invoke(self, request: Request[RequestV, OptionT]) -> ta.Awaitable[Response[ResponseV, OutputT]]:
         return self.service.invoke(request)
 
     @ta.overload
-    def __call__(self, request: Request[RequestV, OptionT]) -> Response[ResponseV, OutputT]:
+    def __call__(self, request: Request[RequestV, OptionT]) -> ta.Awaitable[Response[ResponseV, OutputT]]:
         ...
 
     @ta.overload
-    def __call__(self, v: RequestV, *options: OptionT) -> Response[ResponseV, OutputT]:
+    def __call__(self, v: RequestV, *options: OptionT) -> ta.Awaitable[Response[ResponseV, OutputT]]:
         ...
 
     def __call__(self, o, *args):

ommlds/minichain/services/services.py

@@ -11,4 +11,4 @@ from .responses import ResponseT_co
 
 @ta.runtime_checkable
 class Service(lang.ProtocolForbiddenAsBaseClass, ta.Protocol[RequestT_contra, ResponseT_co]):
-    def invoke(self, request: RequestT_contra) -> ResponseT_co: ...
+    def invoke(self, request: RequestT_contra) -> ta.Awaitable[ResponseT_co]: ...

ommlds/minichain/standard.py

@@ -25,6 +25,14 @@ class Device(tv.UniqueScalarTypedValue[ta.Any], Config):
 ##
 
 
+# TODO: ApiEndpointPath, ApiEndpointUrl, ApiBaseUrl, ...
+class ApiUrl(tv.UniqueScalarTypedValue[str], Config):
+    pass
+
+
+##
+
+
 @dc.dataclass(frozen=True)
 class SecretConfig(Config, lang.Abstract):
     v: sec.SecretRefOrStr = dc.field() | sec.secret_field

ommlds/minichain/stream/services.py

@@ -1,4 +1,6 @@
-import contextlib
+import abc
+import itertools
+import types
 import typing as ta
 
 from omlish import check
@@ -14,6 +16,8 @@ from ..types import Output
 
 
 V = ta.TypeVar('V')
+V2 = ta.TypeVar('V2')
+
 OutputT = ta.TypeVar('OutputT', bound=Output)
 StreamOutputT = ta.TypeVar('StreamOutputT', bound=Output)
 
@@ -31,44 +35,155 @@ StreamOptions: ta.TypeAlias = StreamOption | ResourcesOption
 ##
 
 
-class ResponseGenerator(lang.Final, ta.Generic[V, OutputT]):
-    def __init__(self, g: ta.Generator[V, None, ta.Sequence[OutputT] | None]) -> None:
+class StreamResponseSink(lang.Abstract, ta.Generic[V]):
+    @abc.abstractmethod
+    def emit(self, value: V) -> ta.Awaitable[None]:
+        raise NotImplementedError
+
+
+class StreamResponseIterator(
+    ta.AsyncContextManager['StreamResponseIterator[V, OutputT]'],
+    lang.Abstract,
+    ta.Generic[V, OutputT],
+):
+    @property
+    @abc.abstractmethod
+    def outputs(self) -> tv.TypedValues[OutputT]:
+        raise NotImplementedError
+
+    @ta.final
+    def __aiter__(self) -> ta.AsyncIterator[V]:
+        return self
+
+    @abc.abstractmethod
+    def __anext__(self) -> ta.Awaitable[V]:
+        raise NotImplementedError
+
+
+##
+
+
+class StreamServiceCancelledError(BaseException):
+    pass
+
+
+class StreamServiceNotAwaitedError(Exception):
+    pass
+
+
+class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
+    def __init__(
+            self,
+            fn: ta.Callable[[StreamResponseSink[V]], ta.Awaitable[ta.Sequence[OutputT] | None]],
+    ) -> None:
         super().__init__()
 
-        self._g = g
-        self._is_done = False
+        self._fn = fn
 
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}<{self._g!r}>'
+    @ta.final
+    class _Emit(ta.Generic[V2]):
+        def __init__(self, ssr: '_StreamServiceResponse', value: V2) -> None:
+            self.ssr, self.value = ssr, value
 
-    _outputs: tv.TypedValues[OutputT]
+        done: bool = False
 
-    @property
-    def is_done(self) -> bool:
-        return self._is_done
+        def __await__(self) -> ta.Generator['_StreamServiceResponse._Emit[V2]']:
+            if not self.done:
+                yield self
+            if not self.done:
+                raise StreamServiceNotAwaitedError
+
+    @ta.final
+    class _Sink(StreamResponseSink[V2]):
+        def __init__(self, ssr: '_StreamServiceResponse') -> None:
+            super().__init__()
+
+            self._ssr = ssr
+
+        def emit(self, item: V2) -> ta.Awaitable[None]:
+            return _StreamServiceResponse._Emit(self._ssr, item)
+
+    _state: ta.Literal['new', 'running', 'closed'] = 'new'
+    _sink: _Sink[V]
+    _a: ta.Any
+    _cr: ta.Any
+
+    async def __aenter__(self) -> ta.Self:
+        check.state(self._state == 'new')
+        self._state = 'running'
+        self._sink = _StreamServiceResponse._Sink(self)
+        self._cr = self._fn(self._sink)
+        self._a = self._cr.__await__()
+        self._g = iter(self._a)
+        return self
+
+    @types.coroutine
+    def _aexit(self, exc_type, exc_val, exc_tb):
+        old_state = self._state
+        self._state = 'closed'
+        if old_state != 'running':
+            return
+        if self._cr.cr_running or self._cr.cr_suspended:
+            cex = StreamServiceCancelledError()
+            i = None
+            for n in itertools.count():
+                try:
+                    if not n:
+                        x = self._g.throw(cex)
+                    else:
+                        x = self._g.send(i)
+                except StreamServiceCancelledError as cex2:
+                    if cex2 is cex:
+                        break
+                    raise
+                i = yield x
+        if self._cr.cr_running:
+            raise RuntimeError(f'Coroutine {self._cr!r} not terminated')
+        if self._g is not self._a:
+            self._g.close()
+        self._a.close()
+        self._cr.close()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return await self._aexit(exc_type, exc_val, exc_tb)
+
+    _outputs: tv.TypedValues[OutputT]
 
     @property
     def outputs(self) -> tv.TypedValues[OutputT]:
         return self._outputs
 
-    def __iter__(self) -> ta.Iterator[V]:
-        return self
+    @types.coroutine
+    def _anext(self):
+        check.state(self._state == 'running')
+        i = None
+        while True:
+            try:
+                x = self._g.send(i)
+            except StopIteration as e:
+                if e.value is not None:
+                    self._outputs = tv.TypedValues(*check.isinstance(e.value, ta.Sequence))
+                else:
+                    self._outputs = tv.TypedValues()
+                raise StopAsyncIteration from None
+
+            if isinstance(x, _StreamServiceResponse._Emit) and x.ssr is self:
+                check.state(not x.done)
+                x.done = True
+                return x.value
 
-    def __next__(self) -> V:
-        try:
-            return next(self._g)
-        except StopIteration as e:
-            self._is_done = True
-            if e.value is not None:
-                self._outputs = tv.TypedValues(*check.isinstance(e.value, ta.Sequence))
-            else:
-                self._outputs = tv.TypedValues()
-            raise
+            i = yield x
+
+    async def __anext__(self) -> V:
+        return await self._anext()
+
+
+##
 
 
 StreamResponse: ta.TypeAlias = Response[
     ResourceManaged[
-        ResponseGenerator[
+        StreamResponseIterator[
             V,
             OutputT,
         ],
@@ -77,20 +192,19 @@ StreamResponse: ta.TypeAlias = Response[
 ]
 
 
-def new_stream_response(
+async def new_stream_response(
         rs: Resources,
-        g: ta.Generator[V, None, ta.Sequence[OutputT] | None],
+        fn: ta.Callable[[StreamResponseSink[V]], ta.Awaitable[ta.Sequence[OutputT] | None]],
         outputs: ta.Sequence[StreamOutputT] | None = None,
 ) -> StreamResponse[V, OutputT, StreamOutputT]:
-    return StreamResponse(
-        rs.new_managed(
-            ResponseGenerator(
-                rs.enter_context(
-                    contextlib.closing(
-                        g,
-                    ),
-                ),
-            ),
-        ),
-        outputs or [],
-    )
+    ssr = _StreamServiceResponse(fn)
+
+    v = rs.new_managed(await rs.enter_async_context(ssr))
+    try:
+        return StreamResponse(v, outputs or [])
+    except BaseException:  # noqa
+        # The StreamResponse ctor can raise - for example in `_tv_field_coercer` - in which case we need to clean up the
+        # resources ref we have already allocated before reraising.
+        async with v:
+            pass
+        raise

ommlds/minichain/stream/wrap.py

@@ -1,12 +1,12 @@
 # ruff: noqa: UP028
 import typing as ta
 
-from ..resources import Resources
+from ..resources import UseResources
 from ..services import Request
 from ..services import Service
 from ..types import Output
-from .services import ResponseGenerator
 from .services import StreamResponse
+from .services import StreamResponseSink
 from .services import new_stream_response
 
 
@@ -29,36 +29,34 @@ class WrappedStreamService(ta.Generic[StreamRequestT, V, OutputT, StreamOutputT]
 
     #
 
-    def _process_request(self, request: StreamRequestT) -> StreamRequestT:
+    async def _process_request(self, request: StreamRequestT) -> StreamRequestT:
         return request
 
-    def _process_stream_outputs(self, outputs: ta.Sequence[StreamOutputT]) -> ta.Sequence[StreamOutputT]:
+    async def _process_stream_outputs(self, outputs: ta.Sequence[StreamOutputT]) -> ta.Sequence[StreamOutputT]:
         return outputs
 
-    def _process_vs(self, vs: ta.Iterator[V]) -> ta.Iterator[V]:
-        return vs
+    async def _process_value(self, v: V) -> ta.Iterable[V]:
+        return [v]
 
-    def _process_outputs(self, outputs: ta.Sequence[OutputT]) -> ta.Sequence[OutputT]:
+    async def _process_outputs(self, outputs: ta.Sequence[OutputT]) -> ta.Sequence[OutputT]:
         return outputs
 
     #
 
-    def invoke(self, request: StreamRequestT) -> StreamResponse[V, OutputT, StreamOutputT]:
-        with Resources.new() as rs:
-            in_response = self._inner.invoke(self._process_request(request))
-            in_vs: ResponseGenerator[V, OutputT] = rs.enter_context(in_response.v)
-            out_vs = self._process_vs(in_vs)
-
-            def yield_vs() -> ta.Generator[V, None, ta.Sequence[OutputT] | None]:
-                while True:
-                    try:
-                        out_v = next(out_vs)
-                    except StopIteration as se:
-                        return self._process_outputs(se.value)
-                    yield out_v
-
-            return new_stream_response(
+    async def invoke(self, request: StreamRequestT) -> StreamResponse[V, OutputT, StreamOutputT]:
+        async with UseResources.or_new(request.options) as rs:  # noqa
+            in_resp = await self._inner.invoke(await self._process_request(request))
+            in_vs = await rs.enter_async_context(in_resp.v)
+
+            async def inner(sink: StreamResponseSink[V]) -> ta.Sequence[OutputT] | None:
+                async for in_v in in_vs:
+                    for out_v in (await self._process_value(in_v)):
+                        await sink.emit(out_v)
+
+                return await self._process_outputs(in_vs.outputs)
+
+            return await new_stream_response(
                 rs,
-                yield_vs(),
-                self._process_stream_outputs(in_response.outputs),
+                inner,
+                await self._process_stream_outputs(in_resp.outputs),
             )