diff --git a/tests/test_factory.py b/tests/test_factory.py new file mode 100644 index 0000000..57e997f --- /dev/null +++ b/tests/test_factory.py @@ -0,0 +1,35 @@ +import pytest + +from typeid import TypeID, cached_typeid_factory, typeid_factory +from typeid.errors import PrefixValidationException + + +def test_typeid_factory_generates_typeid_with_prefix(): + gen = typeid_factory("user") + tid = gen() + + assert isinstance(tid, TypeID) + assert tid.prefix == "user" + + +def test_typeid_factory_returns_new_ids_each_time(): + gen = typeid_factory("user") + a = gen() + b = gen() + + assert a != b + + +def test_cached_typeid_factory_is_cached(): + a = cached_typeid_factory("user") + b = cached_typeid_factory("user") + c = cached_typeid_factory("order") + + assert a is b + assert a is not c + + +def test_factory_invalid_prefix_propagates(): + gen = typeid_factory("BAD PREFIX") + with pytest.raises(PrefixValidationException): + gen() diff --git a/typeid/__init__.py b/typeid/__init__.py index 4e1f2d8..b5a0795 100644 --- a/typeid/__init__.py +++ b/typeid/__init__.py @@ -1,3 +1,12 @@ +from .factory import TypeIDFactory, cached_typeid_factory, typeid_factory from .typeid import TypeID, from_string, from_uuid, get_prefix_and_suffix -__all__ = ("TypeID", "from_string", "from_uuid", "get_prefix_and_suffix") +__all__ = ( + "TypeID", + "from_string", + "from_uuid", + "get_prefix_and_suffix", + "TypeIDFactory", + "typeid_factory", + "cached_typeid_factory", +) diff --git a/typeid/factory.py b/typeid/factory.py new file mode 100644 index 0000000..8d907cc --- /dev/null +++ b/typeid/factory.py @@ -0,0 +1,40 @@ +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable + +from .typeid import TypeID + + +@dataclass(frozen=True, slots=True) +class TypeIDFactory: + """ + Callable object that generates TypeIDs with a fixed prefix. + + Example: + user_id = TypeIDFactory("user")() + """ + + prefix: str + + def __call__(self) -> TypeID: + return TypeID(self.prefix) + + +def typeid_factory(prefix: str) -> Callable[[], TypeID]: + """ + Return a zero-argument callable that generates TypeIDs with a fixed prefix. + + Example: + user_id = typeid_factory("user")() + """ + return TypeIDFactory(prefix) + + +@lru_cache(maxsize=256) +def cached_typeid_factory(prefix: str) -> Callable[[], TypeID]: + """ + Same as typeid_factory, but caches factories by prefix. + + Use this if you create factories repeatedly at runtime. + """ + return TypeIDFactory(prefix) diff --git a/typeid/typeid.py b/typeid/typeid.py index 78b16ea..f3d3383 100644 --- a/typeid/typeid.py +++ b/typeid/typeid.py @@ -1,6 +1,6 @@ import uuid import warnings -from typing import Optional +from typing import Generic, Optional, TypeVar import uuid6 @@ -8,24 +8,27 @@ from typeid.errors import InvalidTypeIDStringException from typeid.validation import validate_prefix, validate_suffix +PrefixT = TypeVar("PrefixT", bound=str) -class TypeID: - def __init__(self, prefix: Optional[str] = None, suffix: Optional[str] = None) -> None: + +class TypeID(Generic[PrefixT]): + def __init__(self, prefix: Optional[PrefixT] = None, suffix: Optional[str] = None) -> None: suffix = _convert_uuid_to_b32(uuid6.uuid7()) if not suffix else suffix validate_suffix(suffix=suffix) - if prefix: + + if prefix is not None: validate_prefix(prefix=prefix) - self._prefix = prefix or "" - self._suffix = suffix + self._prefix: Optional[PrefixT] = prefix + self._suffix: str = suffix @classmethod - def from_string(cls, string: str): + def from_string(cls, string: str) -> "TypeID": prefix, suffix = get_prefix_and_suffix(string=string) return cls(suffix=suffix, prefix=prefix) @classmethod - def from_uuid(cls, suffix: uuid.UUID, prefix: Optional[str] = None): + def from_uuid(cls, suffix: uuid.UUID, prefix: Optional[PrefixT] = None) -> "TypeID": suffix_str = _convert_uuid_to_b32(suffix) return cls(suffix=suffix_str, prefix=prefix) @@ -35,7 +38,7 @@ def suffix(self) -> str: @property def prefix(self) -> str: - return self._prefix + return self._prefix or "" @property def uuid(self) -> uuid6.UUID: