From fe2a53860b93d3ba1829938d3451223b87a69010 Mon Sep 17 00:00:00 2001 From: rus07tam Date: Sun, 23 Nov 2025 10:31:23 +0000 Subject: [PATCH] feat: add Marker with object tagging system --- src/snakia/types/__init__.py | 4 ++ src/snakia/types/marker.py | 96 ++++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+) create mode 100644 src/snakia/types/marker.py diff --git a/src/snakia/types/__init__.py b/src/snakia/types/__init__.py index e93cfd2..320dc15 100644 --- a/src/snakia/types/__init__.py +++ b/src/snakia/types/__init__.py @@ -2,12 +2,16 @@ from . import empty from .color import Color +from .marker import Marker, mark, marker from .unique import Unique, UniqueType, unique from .unset import Unset from .version import Version __all__ = [ "Color", + "Marker", + "mark", + "marker", "Version", "UniqueType", "Unique", diff --git a/src/snakia/types/marker.py b/src/snakia/types/marker.py new file mode 100644 index 0000000..002e2ec --- /dev/null +++ b/src/snakia/types/marker.py @@ -0,0 +1,96 @@ +from typing import Any, Callable, Literal, ParamSpec, TypeVar, final, overload + +from typing_extensions import Self + +from snakia.utils import get_attrs, get_or_set_attr + +T = TypeVar("T") +M = TypeVar("M", bound="Marker") +P = ParamSpec("P") + +MARKERS_ATTR = "__snakia_markers__" + + +def _get_all_markers(obj: Any) -> dict[type["Marker"], "Marker"]: + return get_or_set_attr(obj, MARKERS_ATTR, dict[type[Marker], Marker]()) + + +class Marker: + @overload + @classmethod + def get(cls, obj: Any, default: None = None) -> Self: ... + @overload + @classmethod + def get(cls, obj: Any, default: T) -> Self | T: ... + @final + @classmethod + def get(cls, obj: Any, default: Any = None) -> Any: + markers = _get_all_markers(obj) + return markers.get(cls, default) + + @final + @classmethod + def has(cls, obj: Any) -> bool: + if not hasattr(obj, MARKERS_ATTR): + return False + _marker = obj.__dict__[MARKERS_ATTR].get(cls, None) + return isinstance(_marker, cls) + + @final + def set_mark(self, obj: T) -> T: + markers = _get_all_markers(obj) + markers[self.__class__] = self + return obj + + @final + @classmethod + def mark(cls, *a: Any, **kw: Any) -> Callable[[T], T]: + def inner(obj: T) -> T: + return cls(*a, **kw).set_mark(obj) + + return inner + + @final + @classmethod + def unmark(cls, obj: T) -> T: + markers = _get_all_markers(obj) + if cls in markers: + del markers[cls] + return obj + + @overload + @classmethod + def get_marks( + cls, container: Any, *, only_values: Literal[False] = False + ) -> dict[str, tuple[Any, Self]]: ... + @overload + @classmethod + def get_marks( + cls, container: Any, *, only_values: Literal[True] + ) -> dict[str, Self]: ... + @final + @classmethod + def get_marks( + cls, container: Any, *, only_values: bool = False + ) -> dict[str, tuple[Any, Self]] | dict[str, Self]: + markers = {} + for k, v in get_attrs(container).items(): + if not cls.has(v): + continue + _marker = cls.get(v) + if _marker is not None: + markers[k] = _marker if only_values else (v, _marker) + return markers # type: ignore + + +def marker() -> type[Marker]: + return type("", (Marker,), {}) + + +def mark( + func: Callable[P, Marker], /, *args: P.args, **kwargs: P.kwargs +) -> Callable[[T], T]: + def inner(obj: T) -> T: + return func(*args, **kwargs).set_mark(obj) + + return inner