from __future__ import annotations
from datetime import datetime, timezone
from typing import (
TYPE_CHECKING,
Any,
Generic,
TypeVar,
cast,
)
from .. import snowflake
from ..base32 import b32decode, b32encode
from ..entity_types import EntityType
from ..errors import InvalidSkuidError
if TYPE_CHECKING:
from .strings import AnySkuidStr, LongSkuid, ShortSkuid
SEP = "-"
N_SHORT_FORM_PIECES = 2
N_LONG_FORM_PIECES = 4
# 63 bits, must fit in the positive half of an int64.
SNOWFLAKE_UBOUND = 2**63
E = TypeVar("E", bound=EntityType)
# Separate TypeVar for classmethods that must infer their own E from arguments
# rather than from the class parameter (Pylance can't propagate the class-level
# TypeVar when the class is called unparameterised).
_F = TypeVar("_F", bound=EntityType)
GENERIC = EntityType
[docs]
class SnowflakeSkuid(Generic[E]):
"""Representation of a Snowflake-based SKUID."""
entity_type: E
snowflake: int
[docs]
def __init__(self, entity_type: E, snowflake_id: int):
"""
Construct a SnowflakeSkuid.
:param entity_type: The SKA entity type prefix, e.g. ``EntityType.SBD``.
:param snowflake_id: A 63-bit Snowflake integer.
:raises InvalidSkuidError: If *snowflake_id* is outside ``[0, 2**63)``.
"""
if snowflake_id not in range(0, SNOWFLAKE_UBOUND):
raise InvalidSkuidError(
f"Numeric component {snowflake_id} must be in range(0, {SNOWFLAKE_UBOUND})"
)
self.entity_type = cast(E, EntityType.validate(entity_type))
self.snowflake_id: int = snowflake_id
def __str__(self) -> str:
return self.short()
def __repr__(self) -> str:
return f"SnowflakeSkuid({self.entity_type}, {self.snowflake_id})"
def __eq__(self, other: Any) -> bool:
us = (self.entity_type, self.snowflake_id)
them = (
getattr(other, "entity_type", None),
getattr(other, "snowflake_id", None),
)
return us == them
[docs]
@classmethod
def parse(cls, skuid: AnySkuidStr[_F]) -> SnowflakeSkuid[_F]:
"""Parse a SKUID string into a :class:`SnowflakeSkuid` instance.
Accepts both the short form (``type-<base32>``) and the long
form (``type-<gen>-<YYYYMMDD>-<base32>``).
Use sparingly. As a general rule, SKUIDs should be treated as
opaque, atomic identifiers similar to a UUID.
:param skuid: The SKUID string to parse.
:returns: The parsed SKUID.
:raises InvalidSkuidError: If *skuid* is not a valid SKUID string.
"""
try:
pieces = skuid.split(SEP)
if len(pieces) not in (N_SHORT_FORM_PIECES, N_LONG_FORM_PIECES):
raise InvalidSkuidError(f"{skuid} is not a valid short or long form SKUID.")
etype = cast(E, EntityType.validate(pieces[0]))
parsed = cls(etype, b32decode(pieces[-1]))
except (ValueError, AttributeError, TypeError, IndexError) as err:
raise InvalidSkuidError(err) from err
else:
if len(pieces) == N_LONG_FORM_PIECES and parsed.long() != skuid:
# See documentation, the generator and date parts of
# the long-form string must match information
# already encoded in the suffix.
msg = (
f"Invalid long form Snowflake SKUID: suffix {pieces[-1]} "
f"does not encode generator_id={pieces[1]} and/or date={pieces[2]}"
)
raise InvalidSkuidError(msg)
return cast(SnowflakeSkuid[_F], parsed)
@property
def generator_id(self) -> int:
"""The generator ID component of the snowflake."""
return snowflake.extract_generator_id(self.snowflake_id)
@property
def random_suffix(self) -> int:
"""The random suffix component of the snowflake."""
return snowflake.extract_random_suffix(self.snowflake_id)
@property
def timestamp_ms(self) -> int:
"""The timestamp component of the snowflake in milliseconds."""
return snowflake.extract_timestamp_ms(self.snowflake_id)
@property
def datetime(self):
"""The UTC datetime representation of the snowflake timestamp."""
return datetime.fromtimestamp((self.timestamp_ms / 1000), tz=timezone.utc)
[docs]
def short(self) -> ShortSkuid[E]:
"""Return the short string representation of the SKUID."""
base32_id = b32encode(self.snowflake_id)
skuid_str = f"{self.entity_type.value}{SEP}{base32_id}"
return skuid_str # type: ignore[return-value]
[docs]
def long(self) -> LongSkuid[E]:
"""Return the long string representation of the SKUID."""
components = (
self.entity_type.value,
str(self.generator_id),
self.datetime.strftime("%Y%m%d"),
b32encode(self.snowflake_id),
)
skuid_str = SEP.join(components)
return skuid_str # type: ignore[return-value]