Coverage for src/su6/plugins.py: 100%
136 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-17 13:51 +0200
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-17 13:51 +0200
1"""
2Provides a register decorator for third party plugins, and a `include_plugins` (used in cli.py) that loads them.
3"""
4import typing
5from dataclasses import dataclass
6from importlib.metadata import EntryPoint, entry_points
8from rich import print
9from typer import Typer
11from .core import (
12 AbstractConfig,
13 ApplicationState,
14 T_Command,
15 print_json,
16 run_tool,
17 state,
18 with_exit_code,
19)
21__all__ = ["register", "run_tool", "PluginConfig", "print", "print_json"]
24class PluginConfig(AbstractConfig):
25 """
26 Can be inherited in plugin to load in plugin-specific config.
28 The config class is a Singleton, which means multiple instances can be created, but they always have the same state.
30 Example:
31 @register()
32 class DemoConfig(PluginConfig):
33 required_arg: str
34 boolean_arg: bool
35 # ...
38 config = DemoConfig()
40 @register()
41 def with_arguments(required_arg: str, boolean_arg: bool = False) -> None:
42 config.update(required_arg=required_arg, boolean_arg=boolean_arg)
43 print(config)
44 """
46 extras: dict[str, typing.Any]
47 state: typing.Optional[ApplicationState] # only with @register(with_state=True) or after self.attach_state
49 def __init__(self, **kw: typing.Any) -> None:
50 """
51 Initial variables can be passed on instance creation.
52 """
53 super().__init__()
54 self.update(**kw)
55 self.extras = {}
57 def attach_extra(self, name: str, obj: typing.Any) -> None:
58 """
59 Add a non-annotated variable.
60 """
61 self.extras[name] = obj
63 def attach_state(self, global_state: ApplicationState) -> None:
64 """
65 Connect the global state to the plugin config.
66 """
67 self.state = global_state
68 self.attach_extra("state", global_state)
70 def _fields(self) -> typing.Generator[str, typing.Any, None]:
71 yield from self.__annotations__.keys()
72 if self.extras:
73 yield from self.extras.keys() # -> self.state in _values
75 def _get(self, key: str, strict: bool = True) -> typing.Any:
76 notfound = object()
78 value = getattr(self, key, notfound)
79 if value is not notfound:
80 return value
82 if self.extras:
83 value = self.extras.get(key, notfound)
84 if value is not notfound:
85 return value
87 if strict:
88 msg = f"{key} not found in `self {self.__class__.__name__}`"
89 if self.extras:
90 msg += f" or `extra's {self.extras.keys()}`"
91 raise KeyError(msg)
93 def _values(self) -> typing.Generator[typing.Any, typing.Any, None]:
94 yield from (self._get(k, False) for k in self._fields())
96 def __repr__(self) -> str:
97 """
98 Create a readable representation of this class with its data.
100 Stolen from dataclasses._repr_fn.
101 """
102 fields = self._fields()
103 values = self._values()
104 args = ", ".join([f"{f}={v!r}" for f, v in zip(fields, values)])
105 name = self.__class__.__qualname__
106 return f"{name}({args})"
108 def __str__(self) -> str:
109 """
110 Alias for repr.
111 """
112 return repr(self)
115T_PluginConfig = typing.Type[PluginConfig]
117U_Wrappable = typing.Union[T_PluginConfig, T_Command]
118T_Wrappable = typing.TypeVar("T_Wrappable", T_PluginConfig, T_Command)
121@dataclass()
122class Registration(typing.Generic[T_Wrappable]):
123 wrapped: T_Wrappable
125 # Command:
126 add_to_all: bool
127 add_to_fix: bool
129 # Config:
130 with_state: bool
131 strict: bool
132 config_key: typing.Optional[str]
134 args: tuple[typing.Any, ...]
135 kwargs: dict[str, typing.Any]
137 @property
138 def what(self) -> typing.Literal["command", "config"] | None:
139 if isinstance(self.wrapped, type) and issubclass(self.wrapped, PluginConfig):
140 return "config"
141 elif callable(self.wrapped):
142 return "command"
145AnyRegistration = Registration[T_PluginConfig] | Registration[T_Command]
147# WeakValueDictionary() does not work since it removes the references too soon :(
148registrations: dict[int, AnyRegistration] = {}
151def _register(
152 wrapped: T_Wrappable,
153 add_to_all: bool,
154 add_to_fix: bool,
155 with_state: bool,
156 strict: bool,
157 config_key: typing.Optional[str],
158 *a: typing.Any,
159 **kw: typing.Any,
160) -> T_Wrappable:
161 registration = Registration(
162 wrapped,
163 # Command:
164 add_to_all=add_to_all,
165 add_to_fix=add_to_fix,
166 # Config:
167 with_state=with_state,
168 strict=strict,
169 config_key=config_key,
170 # passed to Typer
171 args=a,
172 kwargs=kw,
173 )
175 registrations[id(wrapped)] = registration
176 state.register_plugin(wrapped.__name__, registration)
178 return wrapped
181@typing.overload
182def register(wrappable: T_Wrappable, *a_outer: typing.Any, **kw_outer: typing.Any) -> T_Wrappable:
183 """
184 If wrappable is passed, it returns the same type.
186 @register
187 def func(): ...
189 -> register(func) is called
190 """
193@typing.overload
194def register(
195 wrappable: None = None, *a_outer: typing.Any, **kw_outer: typing.Any
196) -> typing.Callable[[T_Wrappable], T_Wrappable]:
197 """
198 If wrappable is None (empty), it returns a callback that will wrap the function later.
200 @register()
201 def func(): ...
203 -> register() is called
204 """
207def register(
208 wrappable: T_Wrappable = None,
209 # only used when @registering a Command:
210 add_to_all: bool = False,
211 add_to_fix: bool = False,
212 # only used when @registering a PluginConfig:
213 with_state: bool = False,
214 strict: bool = True,
215 config_key: typing.Optional[str] = None,
216 *a_outer: typing.Any,
217 **kw_outer: typing.Any,
218) -> T_Wrappable | typing.Callable[[T_Wrappable], T_Wrappable]:
219 """
220 Register a top-level Plugin command or a Plugin Config.
222 Examples:
223 @register() # () are optional, but you can add Typer keyword arguments if needed.
224 def command():
225 # 'su6 command' is now registered!
226 ...
228 @register() # () are optional, but extra keyword arguments can be passed to configure the config.
229 class MyConfig(PluginConfig):
230 property: str
231 """
233 def inner(func: T_Wrappable) -> T_Wrappable:
234 return _register(func, add_to_all, add_to_fix, with_state, strict, config_key, *a_outer, **kw_outer)
236 if wrappable:
237 return inner(wrappable)
238 else:
239 return inner
242T = typing.TypeVar("T")
245class BoundMethodOf(typing.Protocol[T]):
246 """
247 Protocol to define properties that a bound method has.
248 """
250 __self__: T
251 __name__: str # noqa: A003 - property does exist on the class
252 __doc__: typing.Optional[str] # noqa: A003 - property does exist on the class
254 def __call__(self, a: int) -> str: # pragma: no cover
255 """
256 Indicates this Protocol type can be called.
257 """
260Unbound = typing.Callable[..., typing.Any]
263def unbind(meth: BoundMethodOf[typing.Any] | Unbound) -> typing.Optional[Unbound]:
264 """
265 Extract the original function (which has a different id) from a class method.
266 """
267 return getattr(meth, "__func__", None)
270@dataclass()
271class PluginLoader:
272 app: Typer
273 with_exit_code: bool
275 def main(self) -> None:
276 """
277 Using importlib.metadata, discover available su6 plugins.
279 Example:
280 # pyproject.toml
281 # https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata
282 [project.entry-points."su6"]
283 demo = "su6_plugin_demo.cli" # <- CHANGE ME
284 """
285 discovered_plugins = entry_points(group="su6")
286 for plugin in discovered_plugins: # pragma: nocover
287 self._load_plugin(plugin)
289 self._cleanup()
291 def _cleanup(self) -> None:
292 # registrations.clear()
293 ...
295 def _load_plugin(self, plugin: EntryPoint) -> list[str]:
296 """
297 Look for typer instances and registered commands and configs in an Entrypoint.
299 [project.entry-points."su6"]
300 demo = "su6_plugin_demo.cli"
302 In this case, the entrypoint 'demo' is defined and points to the cli.py module,
303 which gets loaded with `plugin.load()` below.
304 """
305 result = []
306 plugin_module = plugin.load()
308 for item in dir(plugin_module):
309 if item.startswith("_"):
310 continue
312 possible_command = getattr(plugin_module, item)
314 # get method by id (in memory) or first unbind from class and then get by id
315 registration = registrations.get(id(possible_command)) or registrations.get(id(unbind(possible_command)))
317 if isinstance(possible_command, Typer):
318 result += self._add_subcommand(plugin.name, possible_command, plugin_module.__doc__)
319 elif registration and registration.what == "command":
320 result += self._add_command(plugin.name, typing.cast(Registration[T_Command], registration))
321 elif registration and registration.what == "config":
322 result += self._add_config(plugin.name, typing.cast(Registration[T_PluginConfig], registration))
323 # else: ignore
325 return result
327 def _add_command(self, _: str, registration: Registration[T_Command]) -> list[str]:
328 """
329 When a Command Registration is found, it is added to the top-level namespace.
330 """
331 if self.with_exit_code:
332 registration.wrapped = with_exit_code()(registration.wrapped)
333 # adding top-level commands
334 self.app.command(*registration.args, **registration.kwargs)(registration.wrapped)
335 return [f"command {_}"]
337 def _add_config(self, name: str, registration: Registration[T_PluginConfig]) -> list[str]:
338 """
339 When a Config Registration is found, the Singleton data is updated with config from pyproject.toml.
341 Example:
342 # pyproject.toml
343 [tool.su6.demo]
344 boolean-arg = true
345 optional-with-default = "overridden"
347 [tool.su6.demo.extra]
348 more = true
349 """
350 key = registration.config_key or name
352 cls = registration.wrapped
353 inst = cls()
355 if registration.with_state:
356 inst.attach_state(state)
358 if registration.strict is False:
359 inst._strict = False
361 state.attach_plugin_config(key, inst)
362 return [f"config {name}"]
364 def _add_subcommand(self, name: str, subapp: Typer, doc: str) -> list[str]:
365 self.app.add_typer(subapp, name=name, help=doc)
366 return [f"subcommand {name}"]
369def include_plugins(app: Typer, _with_exit_code: bool = True) -> None:
370 """
371 Discover plugins using discover_plugins and add them to either global namespace or as a subcommand.
373 Args:
374 app: the top-level Typer app to append commands to
375 state: top-level application state
376 _with_exit_code: should the @with_exit_code decorator be applied to the return value of the command?
377 """
378 loader = PluginLoader(app, _with_exit_code)
379 loader.main()
382# todo:
383# - add to 'all'
384# - add to 'fix'