omlish 0.0.0.dev451__py3-none-any.whl → 0.0.0.dev453__py3-none-any.whl

omlish/funcs/guard.py

@@ -0,0 +1,214 @@
+import abc
+import functools
+import operator
+import typing as ta
+
+from .. import check
+from .. import collections as col
+from .. import lang
+
+
+T = ta.TypeVar('T')
+T_co = ta.TypeVar('T_co', covariant=True)
+U = ta.TypeVar('U')
+P = ta.ParamSpec('P')
+
+
+##
+
+
+class GuardFn(ta.Protocol[P, T_co]):
+    def __get__(self, instance, owner=None) -> ta.Self: ...
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> ta.Callable[[], T_co] | None: ...
+
+
+##
+
+
+@ta.final
+class DumbGuardFn(ta.Generic[P, T]):
+    def __init__(self, fn: ta.Callable[P, T]) -> None:
+        self._fn = fn
+
+    def __get__(self, instance, owner=None):
+        return DumbGuardFn(self._fn.__get__(instance, owner))  # noqa
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> ta.Callable[[], T]:
+        return functools.partial(self._fn, *args, **kwargs)
+
+
+dumb = DumbGuardFn
+
+
+##
+
+
+class AmbiguousGuardFnError(Exception):
+    pass
+
+
+@ta.final
+class MultiGuardFn(ta.Generic[P, T]):
+    def __init__(
+            self,
+            *children: GuardFn[P, T],
+            strict: bool = False,
+    ) -> None:
+        self._children, self._strict = children, strict
+
+    def __get__(self, instance, owner=None):
+        return MultiGuardFn(*map(operator.methodcaller('__get__', instance, owner), self._children), strict=self._strict)  # noqa
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> ta.Callable[[], T] | None:
+        matches = []
+        for c in self._children:
+            if (m := c(*args, **kwargs)) is not None:
+                if not self._strict:
+                    return m
+                matches.append(m)
+        if not matches:
+            return None
+        elif len(matches) > 1:
+            raise AmbiguousGuardFnError
+        else:
+            return matches[0]
+
+
+multi = MultiGuardFn
+
+
+##
+
+
+class _BaseGuardFnMethod(lang.Abstract, ta.Generic[P, T]):
+    def __init__(
+            self,
+            *,
+            strict: bool = False,
+            requires_override: bool = False,
+            instance_cache: bool = False,
+            prototype: ta.Callable[P, T] | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._strict = strict
+        self._instance_cache = instance_cache
+        self._prototype = prototype
+
+        self._registry: col.AttrRegistry[ta.Callable, None] = col.AttrRegistry(
+            requires_override=requires_override,
+        )
+
+        self._cache: col.AttrRegistryCache[ta.Callable, None, MultiGuardFn] = col.AttrRegistryCache(
+            self._registry,
+            self._prepare,
+        )
+
+    _owner: type | None = None
+    _name: str | None = None
+
+    def __set_name__(self, owner, name):
+        if self._owner is None:
+            self._owner = owner
+        if self._name is None:
+            self._name = name
+
+    def register(self, fn: U) -> U:
+        check.callable(fn)
+        self._registry.register(ta.cast(ta.Callable, fn), None)
+        return fn
+
+    def _prepare(self, instance_cls: type, collected: ta.Mapping[str, tuple[ta.Callable, None]]) -> MultiGuardFn:
+        return MultiGuardFn(
+            *[getattr(instance_cls, a) for a in collected],
+            strict=self._strict,
+        )
+
+    @abc.abstractmethod
+    def _bind(self, instance, owner):
+        raise NotImplementedError
+
+    def __get__(self, instance, owner=None):
+        if instance is None:
+            return self
+
+        if self._instance_cache:
+            try:
+                return instance.__dict__[self._name]
+            except KeyError:
+                pass
+
+        bound = self._bind(instance, owner)
+
+        if self._instance_cache:
+            instance.__dict__[self._name] = bound
+
+        return bound
+
+    def _call(self, *args, **kwargs):
+        instance, *rest = args
+        return self.__get__(instance)(*rest, **kwargs)
+
+#
+
+
+class GuardFnMethod(_BaseGuardFnMethod[P, T]):
+    def _bind(self, instance, owner):
+        return self._cache.get(type(instance)).__get__(instance, owner)  # noqa
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> ta.Callable[[], T] | None:
+        return self._call(*args, **kwargs)
+
+
+def method(
+        *,
+        strict: bool = False,
+        requires_override: bool = False,
+        instance_cache: bool = False,
+) -> ta.Callable[[ta.Callable[P, T]], GuardFnMethod[P, T]]:  # noqa
+    def inner(fn):
+        return GuardFnMethod(
+            strict=strict,
+            requires_override=requires_override,
+            instance_cache=instance_cache,
+            prototype=fn,
+        )
+
+    return inner
+
+
+#
+
+
+class DumbGuardFnMethod(_BaseGuardFnMethod[P, T]):
+    def _bind(self, instance, owner):
+        gf = self._cache.get(type(instance)).__get__(instance, owner)  # noqa
+        x = self._prototype.__get__(instance, owner)  # type: ignore
+
+        def inner(*args, **kwargs):
+            if (m := gf(*args, **kwargs)) is not None:
+                return m()
+            return x(*args, **kwargs)
+
+        return inner
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        return self._call(*args, **kwargs)
+
+
+def dumb_method(
+        *,
+        strict: bool = False,
+        requires_override: bool = False,
+        instance_cache: bool = False,
+) -> ta.Callable[[ta.Callable[P, T]], DumbGuardFnMethod[P, T]]:  # noqa
+    def inner(fn):
+        return DumbGuardFnMethod(
+            strict=strict,
+            requires_override=requires_override,
+            instance_cache=instance_cache,
+            prototype=fn,
+        )
+
+    return inner

omlish/funcs/match.py

@@ -5,6 +5,8 @@ TODO:
  - unify MatchFnClass with dispatch.method?
   - __call__ = mfs.method(); @__call__.register(lambda: ...) def _call_... ?
   - not really the same thing, dispatch is unordered this is necessarily ordered
+  - !! well.. unify interface at least?
+   - guard(*a, **k) -> bool + fn(*a, **k) -> T becomes dispatch(*a, **k) -> (Callable -> T) | None
 """
 import abc
 import dataclasses as dc

omlish/graphs/dot/rendering.py

@@ -29,7 +29,7 @@ class Renderer:
 
         self._out = out
 
-    @dispatch.method
+    @dispatch.method(instance_cache=True)
     def render(self, item: Item) -> None:
         raise TypeError(item)
 

omlish/lite/maybes.py

@@ -1,6 +1,7 @@
 # ruff: noqa: UP007 UP045
 import abc
 import functools
+import operator
 import typing as ta
 
 from .abstract import Abstract
@@ -208,3 +209,10 @@ class _EmptyMaybe(_Maybe[T]):
 
 
 Maybe._empty = _EmptyMaybe()  # noqa
+
+
+##
+
+
+setattr(Maybe, 'just', _JustMaybe)  # noqa
+setattr(Maybe, 'empty', functools.partial(operator.attrgetter('_empty'), Maybe))

omlish/marshal/__init__.py

@@ -113,6 +113,17 @@ with _lang.auto_proxy_init(globals()):
         OptionalUnmarshaler,
     )
 
+    from .composite.unions import (  # noqa
+        MatchUnionMarshaler,
+        MatchUnionUnmarshaler,
+
+        PRIMITIVE_UNION_TYPES,
+        PrimitiveUnionMarshaler,
+        PrimitiveUnionMarshalerFactory,
+        PrimitiveUnionUnmarshaler,
+        PrimitiveUnionUnmarshalerFactory,
+    )
+
     from .composite.wrapped import (  # noqa
         WrappedMarshaler,
         WrappedUnmarshaler,
@@ -197,6 +208,7 @@ with _lang.auto_proxy_init(globals()):
     )
 
     from .polymorphism.metadata import (  # noqa
+        AutoStripSuffix,
         FieldTypeTagging,
         Impl,
         Impls,
@@ -212,13 +224,8 @@ with _lang.auto_proxy_init(globals()):
     )
 
     from .polymorphism.unions import (  # noqa
-        PRIMITIVE_UNION_TYPES,
         PolymorphismUnionMarshalerFactory,
         PolymorphismUnionUnmarshalerFactory,
-        PrimitiveUnionMarshaler,
-        PrimitiveUnionMarshalerFactory,
-        PrimitiveUnionUnmarshaler,
-        PrimitiveUnionUnmarshalerFactory,
     )
 
     from .polymorphism.unmarshal import (  # noqa

omlish/marshal/base/contexts.py

@@ -47,8 +47,9 @@ class MarshalContext(BaseContext, lang.Final):
 
     def make(self, o: ta.Any) -> 'Marshaler':
         rty = self._reflect(o)
+        fac = check.not_none(self.factory)
         try:
-            return check.not_none(self.factory).make_marshaler(self, rty)
+            return fac.make_marshaler(self, rty)
         except mfs.MatchGuardError:
             raise UnhandledTypeError(rty)  # noqa
 
@@ -62,8 +63,9 @@ class UnmarshalContext(BaseContext, lang.Final):
 
     def make(self, o: ta.Any) -> 'Unmarshaler':
         rty = self._reflect(o)
+        fac = check.not_none(self.factory)
         try:
-            return check.not_none(self.factory).make_unmarshaler(self, rty)
+            return fac.make_unmarshaler(self, rty)
         except mfs.MatchGuardError:
             raise UnhandledTypeError(rty)  # noqa
 

omlish/marshal/base/registries.py

@@ -7,7 +7,6 @@ import dataclasses as dc
 import threading
 import typing as ta
 
-from ... import collections as col
 from ... import lang
 
 
@@ -22,18 +21,6 @@ RegistryItemT = ta.TypeVar('RegistryItemT', bound=RegistryItem)
 RegistryItemU = ta.TypeVar('RegistryItemU', bound=RegistryItem)
 
 
-@dc.dataclass(frozen=True)
-class _KeyRegistryItems(ta.Generic[RegistryItemT]):
-    key: ta.Any
-    items: list[RegistryItemT] = dc.field(default_factory=list)
-    item_lists_by_ty: dict[type[RegistryItemT], list[RegistryItemT]] = dc.field(default_factory=dict)
-
-    def add(self, *items: RegistryItemT) -> None:
-        for i in items:
-            self.items.append(i)
-            self.item_lists_by_ty.setdefault(type(i), []).append(i)
-
-
 class RegistrySealedError(Exception):
     pass
 
@@ -50,14 +37,123 @@ class Registry(ta.Generic[RegistryItemT]):
             lock = threading.RLock()
         self._lock = lock
 
-        self._dct: dict[ta.Any, _KeyRegistryItems[RegistryItemT]] = {}
-        self._id_dct: ta.MutableMapping[ta.Any, _KeyRegistryItems[RegistryItemT]] = col.IdentityKeyDict()
+        self._state: Registry._State[RegistryItemT] = Registry._State(
+            dct={},
+            id_dct={},
+            version=0,
+        )
 
-        self._version = 0
         self._sealed = False
 
     #
 
+    @dc.dataclass(frozen=True)
+    class _KeyItems(ta.Generic[RegistryItemU]):
+        key: ta.Any
+        items: ta.Sequence[RegistryItemU] = ()
+        item_lists_by_ty: ta.Mapping[type[RegistryItemU], ta.Sequence[RegistryItemU]] = dc.field(default_factory=dict)
+
+        def add(self, *items: RegistryItemU) -> 'Registry._KeyItems[RegistryItemU]':
+            item_lists_by_ty: dict[type[RegistryItemU], list[RegistryItemU]] = {}
+
+            for i in items:
+                try:
+                    l = item_lists_by_ty[type(i)]
+                except KeyError:
+                    l = item_lists_by_ty[type(i)] = list(self.item_lists_by_ty.get(type(i), ()))
+                l.append(i)
+
+            return Registry._KeyItems(
+                self.key,
+                (*self.items, *items),
+                {**self.item_lists_by_ty, **item_lists_by_ty},
+            )
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    class _State(ta.Generic[RegistryItemU]):
+        dct: ta.Mapping[ta.Any, 'Registry._KeyItems[RegistryItemU]']
+        id_dct: ta.Mapping[ta.Any, 'Registry._KeyItems[RegistryItemU]']
+        version: int
+
+        #
+
+        def register(
+                self,
+                key: ta.Any,
+                *items: RegistryItemT,
+                identity: bool = False,
+        ) -> 'Registry._State[RegistryItemU]':
+            if not items:
+                return self
+
+            sr_dct: ta.Any = self.dct if not identity else self.id_dct
+            if (sr := sr_dct.get(key)) is None:
+                sr = Registry._KeyItems(key)
+            sr = sr.add(*items)
+            new = {key: sr}
+
+            return Registry._State(
+                dct={**self.dct, **new} if not identity else self.dct,
+                id_dct={**self.id_dct, **new} if identity else self.id_dct,
+                version=self.version + 1,
+            )
+
+        #
+
+        _get_cache: dict[ta.Any, ta.Sequence[RegistryItem]] = dc.field(default_factory=dict)
+
+        def get(
+                self,
+                key: ta.Any,
+                *,
+                identity: bool | None = None,
+        ) -> ta.Sequence[RegistryItem]:
+            if identity is None:
+                try:
+                    return self._get_cache[key]
+                except KeyError:
+                    pass
+
+                ret = self._get_cache[key] = (
+                    *self.get(key, identity=True),
+                    *self.get(key, identity=False),
+                )
+                return ret
+
+            dct: ta.Any = self.dct if not identity else self.id_dct
+            try:
+                return dct[key].items
+            except KeyError:
+                return ()
+
+        _get_of_cache: dict[ta.Any, dict[type, ta.Sequence[RegistryItem]]] = dc.field(default_factory=dict)
+
+        def get_of(
+                self,
+                key: ta.Any,
+                item_ty: type[RegistryItem],
+                *,
+                identity: bool | None = None,
+        ) -> ta.Sequence[RegistryItem]:
+            if identity is None:
+                try:
+                    return self._get_of_cache[key][item_ty]
+                except KeyError:
+                    pass
+
+                ret = self._get_of_cache.setdefault(key, {})[item_ty] = (
+                    *self.get_of(key, item_ty, identity=True),
+                    *self.get_of(key, item_ty, identity=False),
+                )
+                return ret
+
+            dct: ta.Any = self.dct if not identity else self.id_dct
+            try:
+                sr = dct[key]
+            except KeyError:
+                return ()
+            return sr.item_lists_by_ty.get(item_ty, ())
+
     def is_sealed(self) -> bool:
         if self._sealed:
             return True
@@ -92,12 +188,11 @@ class Registry(ta.Generic[RegistryItemT]):
             if self._sealed:
                 raise RegistrySealedError(self)
 
-            dct: ta.Any = self._id_dct if identity else self._dct
-            if (sr := dct.get(key)) is None:
-                sr = dct[key] = _KeyRegistryItems(key)
-            sr.add(*items)
-
-            self._version += 1
+            self._state = self._state.register(
+                key,
+                *items,
+                identity=identity,
+            )
 
         return self
 
@@ -109,17 +204,7 @@ class Registry(ta.Generic[RegistryItemT]):
             *,
             identity: bool | None = None,
     ) -> ta.Sequence[RegistryItem]:
-        if identity is None:
-            return (
-                *self.get(key, identity=True),
-                *self.get(key, identity=False),
-            )
-
-        dct: ta.Any = self._id_dct if identity else self._dct
-        try:
-            return dct[key].items
-        except KeyError:
-            return ()
+        return self._state.get(key, identity=identity)
 
     def get_of(
             self,
@@ -128,15 +213,4 @@ class Registry(ta.Generic[RegistryItemT]):
             *,
             identity: bool | None = None,
     ) -> ta.Sequence[RegistryItemU]:
-        if identity is None:
-            return (
-                *self.get_of(key, item_ty, identity=True),
-                *self.get_of(key, item_ty, identity=False),
-            )
-
-        dct: ta.Any = self._id_dct if identity else self._dct
-        try:
-            sr = dct[key]
-        except KeyError:
-            return ()
-        return sr.item_lists_by_ty.get(item_ty, ())
+        return self._state.get_of(key, item_ty, identity=identity)  # type: ignore[return-value]

omlish/marshal/composite/literals.py

@@ -1,3 +1,7 @@
+"""
+TODO:
+ - squash literal unions - typing machinery doesn't
+"""
 import dataclasses as dc
 import typing as ta
 
@@ -26,7 +30,7 @@ class LiteralMarshaler(Marshaler):
 
 class LiteralMarshalerFactory(SimpleMarshalerFactory):
     def guard(self, ctx: MarshalContext, rty: rfl.Type) -> bool:
-        return isinstance(rty, rfl.Literal)
+        return isinstance(rty, rfl.Literal) and len(set(map(type, rty.args))) == 1
 
     def fn(self, ctx: MarshalContext, rty: rfl.Type) -> Marshaler:
         lty = check.isinstance(rty, rfl.Literal)
@@ -45,7 +49,7 @@ class LiteralUnmarshaler(Unmarshaler):
 
 class LiteralUnmarshalerFactory(SimpleUnmarshalerFactory):
     def guard(self, ctx: UnmarshalContext, rty: rfl.Type) -> bool:
-        return isinstance(rty, rfl.Literal)
+        return isinstance(rty, rfl.Literal) and len(set(map(type, rty.args))) == 1
 
     def fn(self, ctx: UnmarshalContext, rty: rfl.Type) -> Unmarshaler:
         lty = check.isinstance(rty, rfl.Literal)