Coverage for src/su6/plugins.py: 100%
136 statements
« prev ^ index » next coverage.py v7.8.1, created at 2025-09-19 22:52 +0200
1"""
2Provides a register decorator for third party plugins, and a `include_plugins` (used in cli.py) that loads them.
3"""
5import typing
6from dataclasses import dataclass
7from importlib.metadata import EntryPoint, entry_points
9from rich import print
10from typer import Typer
12from .core import (
13 AbstractConfig,
14 ApplicationState,
15 T_Command,
16 print_json,
17 run_tool,
18 state,
19 with_exit_code,
20)
22__all__ = ["register", "run_tool", "PluginConfig", "print", "print_json"]
25class PluginConfig(AbstractConfig):
26 """
27 Can be inherited in plugin to load in plugin-specific config.
29 The config class is a Singleton, which means multiple instances can be created, but they always have the same state.
31 Example:
32 @register()
33 class DemoConfig(PluginConfig):
34 required_arg: str
35 boolean_arg: bool
36 # ...
39 config = DemoConfig()
41 @register()
42 def with_arguments(required_arg: str, boolean_arg: bool = False) -> None:
43 config.update(required_arg=required_arg, boolean_arg=boolean_arg)
44 print(config)
45 """
47 extras: dict[str, typing.Any]
48 state: typing.Optional[ApplicationState] # only with @register(with_state=True) or after self.attach_state
50 def __init__(self, **kw: typing.Any) -> None:
51 """
52 Initial variables can be passed on instance creation.
53 """
54 super().__init__()
55 self.update(**kw)
56 self.extras = {}
58 def attach_extra(self, name: str, obj: typing.Any) -> None:
59 """
60 Add a non-annotated variable.
61 """
62 self.extras[name] = obj
64 def attach_state(self, global_state: ApplicationState) -> None:
65 """
66 Connect the global state to the plugin config.
67 """
68 self.state = global_state
69 self.attach_extra("state", global_state)
71 def _fields(self) -> typing.Generator[str, typing.Any, None]:
72 yield from self.__annotations__.keys()
73 if self.extras:
74 yield from self.extras.keys() # -> self.state in _values
76 def _get(self, key: str, strict: bool = True) -> typing.Any:
77 notfound = object()
79 value = getattr(self, key, notfound)
80 if value is not notfound:
81 return value
83 if self.extras:
84 value = self.extras.get(key, notfound)
85 if value is not notfound:
86 return value
88 if strict:
89 msg = f"{key} not found in `self {self.__class__.__name__}`"
90 if self.extras:
91 msg += f" or `extra's {self.extras.keys()}`"
92 raise KeyError(msg)
94 def _values(self) -> typing.Generator[typing.Any, typing.Any, None]:
95 yield from (self._get(k, False) for k in self._fields())
97 def __repr__(self) -> str:
98 """
99 Create a readable representation of this class with its data.
101 Stolen from dataclasses._repr_fn.
102 """
103 fields = self._fields()
104 values = self._values()
105 args = ", ".join([f"{f}={v!r}" for f, v in zip(fields, values)])
106 name = self.__class__.__qualname__
107 return f"{name}({args})"
109 def __str__(self) -> str:
110 """
111 Alias for repr.
112 """
113 return repr(self)
116T_PluginConfig = typing.Type[PluginConfig]
118U_Wrappable = typing.Union[T_PluginConfig, T_Command]
119T_Wrappable = typing.TypeVar("T_Wrappable", T_PluginConfig, T_Command)
122@dataclass()
123class Registration(typing.Generic[T_Wrappable]):
124 wrapped: T_Wrappable
126 # Command:
127 add_to_all: bool
128 add_to_fix: bool
130 # Config:
131 with_state: bool
132 strict: bool
133 config_key: typing.Optional[str]
135 args: tuple[typing.Any, ...]
136 kwargs: dict[str, typing.Any]
138 @property
139 def what(self) -> typing.Literal["command", "config"] | None:
140 if isinstance(self.wrapped, type) and issubclass(self.wrapped, PluginConfig):
141 return "config"
142 elif callable(self.wrapped):
143 return "command"
146AnyRegistration = Registration[T_PluginConfig] | Registration[T_Command]
148# WeakValueDictionary() does not work since it removes the references too soon :(
149registrations: dict[int, AnyRegistration] = {}
152def _register(
153 wrapped: T_Wrappable,
154 add_to_all: bool,
155 add_to_fix: bool,
156 with_state: bool,
157 strict: bool,
158 config_key: typing.Optional[str],
159 *a: typing.Any,
160 **kw: typing.Any,
161) -> T_Wrappable:
162 registration = Registration(
163 wrapped,
164 # Command:
165 add_to_all=add_to_all,
166 add_to_fix=add_to_fix,
167 # Config:
168 with_state=with_state,
169 strict=strict,
170 config_key=config_key,
171 # passed to Typer
172 args=a,
173 kwargs=kw,
174 )
176 registrations[id(wrapped)] = registration
177 state.register_plugin(wrapped.__name__, registration)
179 return wrapped
182@typing.overload
183def register(wrappable: T_Wrappable, *a_outer: typing.Any, **kw_outer: typing.Any) -> T_Wrappable:
184 """
185 If wrappable is passed, it returns the same type.
187 @register
188 def func(): ...
190 -> register(func) is called
191 """
194@typing.overload
195def register(
196 wrappable: None = None, *a_outer: typing.Any, **kw_outer: typing.Any
197) -> typing.Callable[[T_Wrappable], T_Wrappable]:
198 """
199 If wrappable is None (empty), it returns a callback that will wrap the function later.
201 @register()
202 def func(): ...
204 -> register() is called
205 """
208def register(
209 wrappable: T_Wrappable = None,
210 # only used when @registering a Command:
211 add_to_all: bool = False,
212 add_to_fix: bool = False,
213 # only used when @registering a PluginConfig:
214 with_state: bool = False,
215 strict: bool = True,
216 config_key: typing.Optional[str] = None,
217 *a_outer: typing.Any,
218 **kw_outer: typing.Any,
219) -> T_Wrappable | typing.Callable[[T_Wrappable], T_Wrappable]:
220 """
221 Register a top-level Plugin command or a Plugin Config.
223 Examples:
224 @register() # () are optional, but you can add Typer keyword arguments if needed.
225 def command():
226 # 'su6 command' is now registered!
227 ...
229 @register() # () are optional, but extra keyword arguments can be passed to configure the config.
230 class MyConfig(PluginConfig):
231 property: str
232 """
234 def inner(func: T_Wrappable) -> T_Wrappable:
235 return _register(func, add_to_all, add_to_fix, with_state, strict, config_key, *a_outer, **kw_outer)
237 if wrappable:
238 return inner(wrappable)
239 else:
240 return inner
243T = typing.TypeVar("T")
246class BoundMethodOf(typing.Protocol[T]):
247 """
248 Protocol to define properties that a bound method has.
249 """
251 __self__: T
252 __name__: str
253 __doc__: typing.Optional[str]
255 def __call__(self, a: int) -> str: # pragma: no cover
256 """
257 Indicates this Protocol type can be called.
258 """
261Unbound = typing.Callable[..., typing.Any]
264def unbind(meth: BoundMethodOf[typing.Any] | Unbound) -> typing.Optional[Unbound]:
265 """
266 Extract the original function (which has a different id) from a class method.
267 """
268 return getattr(meth, "__func__", None)
271@dataclass()
272class PluginLoader:
273 app: Typer
274 with_exit_code: bool
276 def main(self) -> None:
277 """
278 Using importlib.metadata, discover available su6 plugins.
280 Example:
281 # pyproject.toml
282 # https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata
283 [project.entry-points."su6"]
284 demo = "su6_plugin_demo.cli" # <- CHANGE ME
285 """
286 discovered_plugins = entry_points(group="su6")
287 for plugin in discovered_plugins: # pragma: nocover
288 self._load_plugin(plugin)
290 self._cleanup()
292 def _cleanup(self) -> None:
293 # registrations.clear()
294 ...
296 def _load_plugin(self, plugin: EntryPoint) -> list[str]:
297 """
298 Look for typer instances and registered commands and configs in an Entrypoint.
300 [project.entry-points."su6"]
301 demo = "su6_plugin_demo.cli"
303 In this case, the entrypoint 'demo' is defined and points to the cli.py module,
304 which gets loaded with `plugin.load()` below.
305 """
306 result = []
307 plugin_module = plugin.load()
309 for item in dir(plugin_module):
310 if item.startswith("_"):
311 continue
313 possible_command = getattr(plugin_module, item)
315 # get method by id (in memory) or first unbind from class and then get by id
316 registration = registrations.get(id(possible_command)) or registrations.get(id(unbind(possible_command)))
318 if isinstance(possible_command, Typer):
319 result += self._add_subcommand(plugin.name, possible_command, plugin_module.__doc__)
320 elif registration and registration.what == "command":
321 result += self._add_command(plugin.name, typing.cast(Registration[T_Command], registration))
322 elif registration and registration.what == "config":
323 result += self._add_config(plugin.name, typing.cast(Registration[T_PluginConfig], registration))
324 # else: ignore
326 return result
328 def _add_command(self, _: str, registration: Registration[T_Command]) -> list[str]:
329 """
330 When a Command Registration is found, it is added to the top-level namespace.
331 """
332 if self.with_exit_code:
333 registration.wrapped = with_exit_code()(registration.wrapped)
334 # adding top-level commands
335 self.app.command(*registration.args, **registration.kwargs)(registration.wrapped)
336 return [f"command {_}"]
338 def _add_config(self, name: str, registration: Registration[T_PluginConfig]) -> list[str]:
339 """
340 When a Config Registration is found, the Singleton data is updated with config from pyproject.toml.
342 Example:
343 # pyproject.toml
344 [tool.su6.demo]
345 boolean-arg = true
346 optional-with-default = "overridden"
348 [tool.su6.demo.extra]
349 more = true
350 """
351 key = registration.config_key or name
353 cls = registration.wrapped
354 inst = cls()
356 if registration.with_state:
357 inst.attach_state(state)
359 if registration.strict is False:
360 inst._strict = False
362 state.attach_plugin_config(key, inst)
363 return [f"config {name}"]
365 def _add_subcommand(self, name: str, subapp: Typer, doc: str) -> list[str]:
366 self.app.add_typer(subapp, name=name, help=doc)
367 return [f"subcommand {name}"]
370def include_plugins(app: Typer, _with_exit_code: bool = True) -> None:
371 """
372 Discover plugins using discover_plugins and add them to either global namespace or as a subcommand.
374 Args:
375 app: the top-level Typer app to append commands to
376 state: top-level application state
377 _with_exit_code: should the @with_exit_code decorator be applied to the return value of the command?
378 """
379 loader = PluginLoader(app, _with_exit_code)
380 loader.main()
383# todo:
384# - add to 'all'
385# - add to 'fix'