Source code for jitx.decorators

from collections.abc import Callable
from dataclasses import dataclass
from typing import dataclass_transform, overload

from jitx._structural import PrePostInit




def early(method: Callable):
    """Call this method before calling other initializers in this class."""
    return PrePostInit(
        lambda inst, _: method.__get__(inst, type(inst))(),
        lambda _, before: before,
        representing=method,
    )





def late(method: Callable):
    """Call this method after calling all other initializers in this class."""
    return PrePostInit(
        lambda _, __: None,
        lambda inst, before: method.__get__(inst, type(inst))(),
        representing=method,
    )



@overload
def identityclass[T](cls: type[T]) -> type[T]: ...


@overload
def identityclass[T](**kwargs) -> Callable[[type[T]], type[T]]: ...




@dataclass_transform()
def identityclass[T](cls: type[T] | None = None, **kwargs):
    """Decorator to make a dataclass but setting eq and hash to object
    identity. This allows a dataclass to serve as a structural jitx type."""
    if cls is None:
        return lambda cls: identityclass(cls, **kwargs)
    cls.__eq__ = lambda self, other: self is other
    cls.__hash__ = lambda self: hash(id(self))
    cls.__jitx_identityclass__ = True  # type: ignore (yes it doesn't exist, we're adding it)
    return dataclass(**kwargs)(cls)