Python TypeID implementation

typeid.py

from collections.abc import Sequence
from functools import cached_property
from typing import Any
from typing import ClassVar
from typing import LiteralString
from uuid import UUID

from uuid_utils import uuid7

ALPHABET = "0123456789abcdefghjkmnpqrstvwxyz"
INDEX_MAP = {char: i for i, char in enumerate(ALPHABET)}
BASE32_ENCODED_UUID_LENGTH = 26


def encode(value: Sequence[int]) -> str:
    bits = "00" + "".join(f"{byte:08b}" for byte in value)
    encoded = "".join(ALPHABET[int(bits[i : i + 5], 2)] for i in range(0, len(bits), 5))
    if len(encoded) != BASE32_ENCODED_UUID_LENGTH:
        raise ValueError("Encoded string length does not match expected 26 characters.")

    return encoded


def decode(value: str) -> bytes:
    if len(value) != BASE32_ENCODED_UUID_LENGTH:
        raise ValueError("Encoded string must be exactly 26 characters long.")

    bits = 0
    for char in value:
        try:
            bits = (bits << 5) | INDEX_MAP[char]
        except KeyError:
            raise ValueError(f"Invalid character {char!r} in encoded string.")

    return bits.to_bytes(16, byteorder="big")


class TypeID(UUID):
    """
    Usage:
        class UserID(TypeID):
            prefix = "user"

        UserID()  # will generate a new TypeID with 'user' prefix
        UserID('user_2x4y6z8a0b1c2d3e4f5g6h7j8k')  # will make sure prefix matches
        UserID('5d278df4-280b-0b04-d1b8-8f2c0d13c913')  # will produce the same instance as above

    Format:
         user_2x4y6z8a0b1c2d3e4f5g6h7j8k
         └──┘ └────────────────────────┘
        prefix   uuid suffix (base32)

    Spec:
        https://github.com/jetpack-io/typeid/tree/main/spec
    """

    prefix: ClassVar[LiteralString]
      
    def __init_subclass__(cls, *args, **kwargs) -> None:
        super()__init_subclass__(*args, **kwargs)
        if not hasattr(cls, "prefix"):
            raise TypeError(f"Subclass of TypeID {cls.__name__} must define a prefix")

    def __init__(self, hex: str | None = None, *args: Any, **kwargs: Any) -> None:  # noqa: A002
        """
        Same as UUID, except:
        - can't be instantiated without prefix defined in a subclass
        - if no parameters were given, new typeid is generated
        - first parameter additionally may be a string representation of TypeID
        """
        if not hasattr(self, "prefix"):
            raise TypeError("Cannot use TypeID directly. Define a subclass with a prefix")

        if hex is not None and not args and not kwargs:
            try:
                prefix, value = hex.split("_")
            except ValueError:  # not a canonical TypeID string, expecting a valid UUID hex
                pass
            else:
                if prefix != self.prefix:
                    raise TypeError(
                        f"Prefix mismatch for {self.__class__}: "
                        f"expected {self.prefix!r}, got {prefix}"
                    )
                hex = None
                kwargs["bytes"] = decode(value)
        elif not args and not kwargs:
            kwargs["bytes"] = uuid7().bytes

        super().__init__(hex, *args, **kwargs)

    @cached_property
    def suffix(self) -> str:
        return encode(self.bytes)

    @cached_property
    def uuid(self) -> str:
        return (
            f"{self.hex[:8]}-{self.hex[8:12]}-{self.hex[12:16]}-{self.hex[16:20]}-{self.hex[20:]}"
        )

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}('{self}')"

    def __str__(self) -> str:
        return f"{self.prefix}_{self.suffix}"