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