Coverage for src / cufile_patcher / registry.py: 100%

1from __future__ import annotations 

2 

3from collections.abc import Callable 

4 

5from .plugins import CuFileBackend, SystemCuFileBackend 

6 

7 

8class BackendRegistry: 

9 def __init__(self) -> None: 

10 self._factories: dict[str, Callable[[], CuFileBackend]] = {} 

11 self._instances: dict[str, CuFileBackend] = {} 

12 self._default_name = "system" 

13 

14 def register(self, name: str, factory: Callable[[], CuFileBackend]) -> None: 

15 if not name: 

16 raise ValueError("backend name must not be empty") 

17 self._factories[name] = factory 

18 self._instances.pop(name, None) 

19 

20 def set_default(self, name: str) -> None: 

21 if name not in self._factories: 

22 raise KeyError(f"backend '{name}' is not registered") 

23 self._default_name = name 

24 

25 def get_default(self) -> CuFileBackend: 

26 return self.get(self._default_name) 

27 

28 def get(self, name: str) -> CuFileBackend: 

29 if name in self._instances: 

30 return self._instances[name] 

31 try: 

32 factory = self._factories[name] 

33 except KeyError as exc: 

34 raise KeyError(f"backend '{name}' is not registered") from exc 

35 backend = factory() 

36 self._instances[name] = backend 

37 return backend 

38 

39 

40_registry = BackendRegistry() 

41_registry.register("system", SystemCuFileBackend) 

42 

43 

44def register_backend(name: str, factory: Callable[[], CuFileBackend]) -> None: 

45 _registry.register(name, factory) 

46 

47 

48def set_default_backend(name: str) -> None: 

49 _registry.set_default(name) 

50 

51 

52def get_default_backend() -> CuFileBackend: 

53 return _registry.get_default()