Coverage for /usr/local/lib/python3.12/site-packages/prefect/utilities/dispatch.py: 66%
72 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 11:21 +0000
1"""
2Provides methods for performing dynamic dispatch for actions on base type to one of its
3subtypes.
5Example:
7```python
8@register_base_type
9class Base:
10 @classmethod
11 def __dispatch_key__(cls):
12 return cls.__name__.lower()
15class Foo(Base):
16 ...
18key = get_dispatch_key(Foo) # 'foo'
19lookup_type(Base, key) # Foo
20```
21"""
23import abc 1a
24import inspect 1a
25import warnings 1a
26from typing import Any, Literal, Optional, TypeVar, overload 1a
28T = TypeVar("T", bound=type[Any]) 1a
30_TYPE_REGISTRIES: dict[Any, dict[str, Any]] = {} 1a
33def get_registry_for_type(cls: T) -> Optional[dict[str, T]]: 1a
34 """
35 Get the first matching registry for a class or any of its base classes.
37 If not found, `None` is returned.
38 """
39 return next( 1ab
40 (reg for cls in cls.mro() if (reg := _TYPE_REGISTRIES.get(cls)) is not None),
41 None,
42 )
45@overload 1a
46def get_dispatch_key( 46 ↛ exitline 46 didn't return from function 'get_dispatch_key' because 1a
47 cls_or_instance: Any, allow_missing: Literal[False] = False
48) -> str: ...
51@overload 1a
52def get_dispatch_key( 52 ↛ exitline 52 didn't return from function 'get_dispatch_key' because 1a
53 cls_or_instance: Any, allow_missing: Literal[True] = ...
54) -> Optional[str]: ...
57def get_dispatch_key( 1a
58 cls_or_instance: Any, allow_missing: bool = False
59) -> Optional[str]:
60 """
61 Retrieve the unique dispatch key for a class type or instance.
63 This key is defined at the `__dispatch_key__` attribute. If it is a callable, it
64 will be resolved.
66 If `allow_missing` is `False`, an exception will be raised if the attribute is not
67 defined or the key is null. If `True`, `None` will be returned in these cases.
68 """
69 dispatch_key = getattr(cls_or_instance, "__dispatch_key__", None) 1ab
71 type_name = ( 1ab
72 cls_or_instance.__name__
73 if isinstance(cls_or_instance, type)
74 else type(cls_or_instance).__name__
75 )
77 if dispatch_key is None: 77 ↛ 78line 77 didn't jump to line 78 because the condition on line 77 was never true1ab
78 if allow_missing:
79 return None
80 raise ValueError(
81 f"Type {type_name!r} does not define a value for "
82 "'__dispatch_key__' which is required for registry lookup."
83 )
85 if callable(dispatch_key): 85 ↛ 88line 85 didn't jump to line 88 because the condition on line 85 was always true1ab
86 dispatch_key = dispatch_key() 1ab
88 if allow_missing and dispatch_key is None: 1ab
89 return None 1a
91 if not isinstance(dispatch_key, str): 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true1ab
92 raise TypeError(
93 f"Type {type_name!r} has a '__dispatch_key__' of type "
94 f"{type(dispatch_key).__name__} but a type of 'str' is required."
95 )
97 return dispatch_key 1ab
100@classmethod 1a
101def _register_subclass_of_base_type(cls: type[Any], **kwargs: Any) -> None: 1a
102 if hasattr(cls, "__init_subclass_original__"): 1ab
103 cls.__init_subclass_original__(**kwargs) 1a
104 elif hasattr(cls, "__pydantic_init_subclass_original__"): 104 ↛ 108line 104 didn't jump to line 108 because the condition on line 104 was always true1ab
105 cls.__pydantic_init_subclass_original__(**kwargs) 1ab
107 # Do not register abstract base classes
108 if abc.ABC in cls.__bases__: 1ab
109 return 1a
111 register_type(cls) 1ab
114def register_base_type(cls: T) -> T: 1a
115 """
116 Register a base type allowing child types to be registered for dispatch with
117 `register_type`.
119 The base class may or may not define a `__dispatch_key__` to allow lookups of the
120 base type.
121 """
122 registry = _TYPE_REGISTRIES.setdefault(cls, {}) 1a
123 base_key = get_dispatch_key(cls, allow_missing=True) 1a
124 if base_key is not None: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true1a
125 registry[base_key] = cls
127 # Add automatic subtype registration
128 if hasattr(cls, "__pydantic_init_subclass__"): 1a
129 cls.__pydantic_init_subclass_original__ = getattr( 1a
130 cls, "__pydantic_init_subclass__"
131 )
132 cls.__pydantic_init_subclass__ = _register_subclass_of_base_type 1a
133 else:
134 cls.__init_subclass_original__ = getattr(cls, "__init_subclass__") 1a
135 setattr(cls, "__init_subclass__", _register_subclass_of_base_type) 1a
137 return cls 1a
140def register_type(cls: T) -> T: 1a
141 """
142 Register a type for lookup with dispatch.
144 The type or one of its parents must define a unique `__dispatch_key__`.
146 One of the classes base types must be registered using `register_base_type`.
147 """
148 # Lookup the registry for this type
149 registry = get_registry_for_type(cls) 1ab
151 # Check if a base type is registered
152 if registry is None: 152 ↛ 154line 152 didn't jump to line 154 because the condition on line 152 was never true1ab
153 # Include a description of registered base types
154 known = ", ".join(repr(base.__name__) for base in _TYPE_REGISTRIES)
155 known_message = (
156 f" Did you mean to inherit from one of the following known types: {known}."
157 if known
158 else ""
159 )
161 # And a list of all base types for the type they tried to register
162 bases = ", ".join(
163 repr(base.__name__) for base in cls.mro() if base not in (object, cls)
164 )
166 raise ValueError(
167 f"No registry found for type {cls.__name__!r} with bases {bases}."
168 + known_message
169 )
171 key = get_dispatch_key(cls) 1ab
172 existing_value = registry.get(key) 1ab
173 if existing_value is not None and id(existing_value) != id(cls): 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true1ab
174 try:
175 # Get line numbers for debugging
176 file = inspect.getsourcefile(cls)
177 line_number = inspect.getsourcelines(cls)[1]
178 existing_file = inspect.getsourcefile(existing_value)
179 existing_line_number = inspect.getsourcelines(existing_value)[1]
180 warnings.warn(
181 f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that "
182 f"matches existing registered type {existing_value.__name__!r} from "
183 f"{existing_file}:{existing_line_number}. The existing type will be "
184 "overridden."
185 )
186 except OSError:
187 # If we can't get the source, another actor is loading this class via eval
188 # and we shouldn't update the registry
189 return cls
191 # Add to the registry
192 registry[key] = cls 1ab
194 return cls 1ab
197def lookup_type(cls: T, dispatch_key: str) -> T: 1a
198 """
199 Look up a dispatch key in the type registry for the given class.
200 """
201 # Get the first matching registry for the class or one of its bases
202 registry = get_registry_for_type(cls) or {}
204 # Look up this type in the registry
205 subcls = registry.get(dispatch_key)
207 if subcls is None:
208 raise KeyError(
209 f"No class found for dispatch key {dispatch_key!r} in registry for type "
210 f"{cls.__name__!r}."
211 )
213 return subcls