Simple backend dispatch for collection-creation functions
| 51 | |
| 52 | |
| 53 | class CreationDispatch(Generic[BackendEntrypointType]): |
| 54 | """Simple backend dispatch for collection-creation functions""" |
| 55 | |
| 56 | _lookup: dict[str, BackendEntrypointType] |
| 57 | _module_name: str |
| 58 | _config_field: str |
| 59 | _default: str |
| 60 | _entrypoint_class: type[BackendEntrypointType] |
| 61 | _entrypoint_root: str |
| 62 | |
| 63 | def __init__( |
| 64 | self, |
| 65 | module_name: str, |
| 66 | default: str, |
| 67 | entrypoint_class: type[BackendEntrypointType], |
| 68 | name: str | None = None, |
| 69 | entrypoint_root: str = "dask", |
| 70 | ): |
| 71 | self._lookup = {} |
| 72 | self._module_name = module_name |
| 73 | self._config_field = f"{module_name}.backend" |
| 74 | self._default = default |
| 75 | self._entrypoint_class = entrypoint_class |
| 76 | self._entrypoint_root = entrypoint_root |
| 77 | if name: |
| 78 | self.__name__ = name |
| 79 | |
| 80 | def register_backend( |
| 81 | self, name: str, backend: BackendEntrypointType |
| 82 | ) -> BackendEntrypointType: |
| 83 | """Register a target class for a specific array-backend label""" |
| 84 | if not isinstance(backend, self._entrypoint_class): |
| 85 | raise ValueError( |
| 86 | f"This CreationDispatch only supports " |
| 87 | f"{self._entrypoint_class} registration. " |
| 88 | f"Got {type(backend)}" |
| 89 | ) |
| 90 | self._lookup[name] = backend |
| 91 | return backend |
| 92 | |
| 93 | def dispatch(self, backend: str): |
| 94 | """Return the desired backend entrypoint""" |
| 95 | try: |
| 96 | impl = self._lookup[backend] |
| 97 | except KeyError: |
| 98 | # Check entrypoints for the specified backend |
| 99 | entrypoints = detect_entrypoints( |
| 100 | f"{self._entrypoint_root}.{self._module_name}.backends" |
| 101 | ) |
| 102 | if backend in entrypoints: |
| 103 | return self.register_backend(backend, entrypoints[backend].load()()) |
| 104 | else: |
| 105 | return impl |
| 106 | raise ValueError(f"No backend dispatch registered for {backend}") |
| 107 | |
| 108 | @property |
| 109 | def backend(self) -> str: |
| 110 | """Return the desired collection backend""" |
no outgoing calls
no test coverage detected