| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any, TypeVar |
| 4 | |
| 5 | _T = TypeVar("_T", bound=type) |
| 6 | |
| 7 | |
| 8 | class ExtensionManager: |
| 9 | """A registry for managing pluggable extension classes. |
| 10 | |
| 11 | This class provides a simple mechanism to register and instantiate |
| 12 | extension classes by name. It is commonly used to implement plugin |
| 13 | systems where different implementations can be swapped at runtime. |
| 14 | |
| 15 | Examples: |
| 16 | Basic usage with a registry instance: |
| 17 | |
| 18 | >>> FOO_REGISTRY = ExtensionManager() |
| 19 | >>> @FOO_REGISTRY.register("my_foo_impl") |
| 20 | ... class MyFooImpl(Foo): |
| 21 | ... def __init__(self, value): |
| 22 | ... self.value = value |
| 23 | >>> foo_impl = FOO_REGISTRY.load("my_foo_impl", value=123) |
| 24 | |
| 25 | """ |
| 26 | |
| 27 | def __init__(self) -> None: |
| 28 | """Initialize an empty extension registry.""" |
| 29 | self.name2class: dict[str, type] = {} |
| 30 | |
| 31 | def register(self, name: str): |
| 32 | """Decorator to register a class with the given name.""" |
| 33 | |
| 34 | def wrap(cls_to_register: _T) -> _T: |
| 35 | self.name2class[name] = cls_to_register |
| 36 | return cls_to_register |
| 37 | |
| 38 | return wrap |
| 39 | |
| 40 | def load(self, cls_name: str, *args, **kwargs) -> Any: |
| 41 | """Instantiate and return a registered extension class by name.""" |
| 42 | cls = self.name2class.get(cls_name) |
| 43 | assert cls is not None, f"Extension class {cls_name} not found" |
| 44 | return cls(*args, **kwargs) |