1from collections.abc import Callable
 2from typing import TYPE_CHECKING, Any, overload
 3
 4from .parameters import JobParameters
 5
 6if TYPE_CHECKING:
 7    from .jobs import Job
 8
 9
10class JobsRegistry:
11    def __init__(self) -> None:
12        self.jobs: dict[str, type[Job]] = {}
13        self.ready = False
14
15    def register_job(self, job_class: type[Job], alias: str = "") -> None:
16        name = self.get_job_class_name(job_class)
17        self.jobs[name] = job_class
18
19        if alias:
20            self.jobs[alias] = job_class
21
22    def get_job_class_name(self, job_class: type[Job]) -> str:
23        return f"{job_class.__module__}.{job_class.__qualname__}"
24
25    def get_job_class(self, name: str) -> type[Job]:
26        return self.jobs[name]
27
28    def load_job(
29        self, job_class_name: str, parameters: dict[str, Any] | None = None
30    ) -> Job:
31        if not self.ready:
32            raise RuntimeError("Jobs registry is not ready yet")
33
34        job_class = self.get_job_class(job_class_name)
35        args, kwargs = JobParameters.from_json(parameters or {})
36        return job_class(*args, **kwargs)
37
38
39jobs_registry = JobsRegistry()
40
41
42@overload
43def register_job[T: type["Job"]](job_class: T, /, *, alias: str = ...) -> T: ...
44@overload
45def register_job[T: type["Job"]](*, alias: str = ...) -> Callable[[T], T]: ...
46def register_job[T: type["Job"]](
47    job_class: T | None = None, *, alias: str = ""
48) -> T | Callable[[T], T]:
49    """
50    A decorator that registers a job class in the jobs registry with an optional alias.
51    Can be used both with and without parentheses.
52    """
53    if job_class is None:
54
55        def wrapper(cls: T) -> T:
56            jobs_registry.register_job(cls, alias=alias)
57            return cls
58
59        return wrapper
60    else:
61        jobs_registry.register_job(job_class, alias=alias)
62        return job_class