feat: add Marker with object tagging system

This commit is contained in:
rus07tam 2025-11-23 10:31:23 +00:00
parent 346140d675
commit fe2a53860b
2 changed files with 100 additions and 0 deletions

View file

@ -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",

View file

@ -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("<anonym marker>", (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