from __future__ import annotations
import itertools
import random
import sys
from datetime import UTC, datetime
from enum import StrEnum
from typing import Annotated as A
from typing import Any, Generic, Literal, TypeAlias
from pydantic import AfterValidator, Discriminator, Tag, TypeAdapter
from noob.const import META_SIGNAL
from noob.types import Epoch, Picklable, SerializableDatetime
if sys.version_info < (3, 12):
from typing_extensions import TypedDict, TypeVar
elif sys.version_info < (3, 13):
from typing import TypedDict
from typing_extensions import TypeVar
else:
from typing import TypedDict, TypeVar
_TEvent = TypeVar("_TEvent", default=Any)
_ID_SEQUENCE_BITS = 32
"""
Bits reserved for the id sequence - the per-event identifier.
the random base occupies the bits above it.
"""
_ID_PREFIX_BITS = 32
"""
Bits of randomness in an :class:`.EventMaker` base.
This is probably overkill, but keeping it high-ish for now to leave some space for networked runners
"""
[docs]
class Event(TypedDict, Generic[_TEvent]):
"""
Container for a single value returned from a single :meth:`.Node.process` call
"""
id: int
"""
An integer event ID that combines some random prefix with a monotonic sequence counter.
Constructed like::
(base << 32) ^ sequence
The base and sequence should both be integers < 2**32,
but for performance/YAGNI reasons bounds are not checked (for now).
EventIDs should be treated as opaque identifiers for now,
they bear no inherent information about an event source
(except that events were produced within the same context).
In the future the ID system may be unified such that events,
nodes, runners, etc. have a unified hierarchical identifier system.
See Also:
:func:`.event_id` , :func:`.event_id_parts`
"""
timestamp: SerializableDatetime
"""Timestamp of when the event was received by the :class:`.TubeRunner`"""
node_id: str
"""ID of node that emitted the event"""
signal: str
"""name of the signal that emitted the event"""
epoch: Epoch
"""Epoch number the event was emitted in"""
value: Picklable[_TEvent]
"""Value emitted by the processing node"""
[docs]
def is_event(instance: Any) -> bool:
"""
TypedDicts don't support instancechecks by default,
we want to use typed dicts for perf's sake,
but we still also would like to check if something is an event
"""
if not isinstance(instance, dict):
return False
return all(key in instance for key in Event.__annotations__)
[docs]
def event_id(prefix: int, sequence: int, shifted: bool = False) -> int:
"""
Construct an event id from a prefix and a sequence like::
(prefix << _ID_SEQUENCE_BITS) ^ sequence
Args:
prefix (int): A 32-bit prefix
sequence (int): A 32-bit monotonically increasing sequence
shifted (bool): If ``True`` , prefix is already shifted
(for performance's sake, avoid an unnecessary operation on a high-call path)
Returns:
A 64-bit event ID
"""
if shifted:
return prefix ^ sequence
else:
return (prefix << _ID_SEQUENCE_BITS) ^ sequence
[docs]
def event_id_parts(event_id: int) -> tuple[int, int]:
"""
Decompose an event ID into its prefix and sequence identifier
Args:
event_id (int): A 64-bit event ID
Returns:
tuple[int, int] - a tuple of the ``(prefix, sequence)``
"""
prefix = event_id >> _ID_SEQUENCE_BITS
sequence = event_id & ((1 << _ID_SEQUENCE_BITS) - 1)
return (prefix, sequence)
[docs]
class EventMaker:
"""
Convenience constructor for :class:`.Event` s / :class:`.MetaEvent` s for a single execution
context (a scheduler, store, node, ...), stamping each with a (locally) unique id and timestamp.
Prefixes are created with :mod:`random` rather than :mod:`secrets` ,
because they are only used for local disambiguation for now.
IDs will be made secure for signing events
when we need to handle networked event processing
"""
__slots__ = ("prefix", "node_id", "_sequence", "_shifted")
def __init__(self, prefix: int | None = None, node_id: str | None = None) -> None:
self.prefix = random.getrandbits(_ID_PREFIX_BITS) if prefix is None else prefix
self._shifted = self.prefix << _ID_SEQUENCE_BITS
self.node_id = node_id
self._sequence = itertools.count()
[docs]
def new_event(
self,
value: _TEvent,
epoch: Epoch,
node_id: str | None = None,
signal: str = "value",
timestamp: SerializableDatetime | None = None,
) -> Event[_TEvent]:
"""
Make an :class:`.Event`, filling in the id and (UTC) timestamp.
Pass ``timestamp`` to share one across a batch of events emitted together.
"""
node_id = node_id or self.node_id
if not node_id:
raise ValueError(
"node_id not set on EventMaker instance, and none provided when creating event"
)
return Event(
id=event_id(self._shifted, next(self._sequence), shifted=True),
timestamp=datetime.now(UTC) if timestamp is None else timestamp,
node_id=node_id,
signal=signal,
epoch=epoch,
value=value,
)
def _type_discriminator(v: dict | Event | MetaEvent) -> str:
if v.get("node_id", None) == "meta":
return "meta"
else:
return "event"
def _meta_signals_to_enum(evt: Event) -> Event:
"""If we are a meta signal, ensure the value is cast to the enum rather than string value"""
if (
isinstance(evt["value"], str)
and evt["value"].startswith(META_SIGNAL)
and (metaevt_key := evt["value"].replace(META_SIGNAL, "")) in MetaSignal.__members__
):
evt["value"] = MetaSignal.__members__[metaevt_key]
return evt
EventUnion = A[
A[
Event,
AfterValidator(_meta_signals_to_enum),
Tag("event"),
]
| A[MetaEvent, Tag("meta")],
Discriminator(_type_discriminator),
]
EventAdapter = TypeAdapter[EventUnion](EventUnion)
_NoEventableInner = TypeVar("_NoEventableInner")
NoEventable: TypeAlias = (
_NoEventableInner | "Event[_NoEventableInner]" | Literal[MetaSignal.NoEvent]
)
"""Convenience generic type to indicate that some signal can be a NoEvent"""