Coverage for src/su6/plugins.py: 100%

1"""

2 Provides a register decorator for third party plugins, and a `include_plugins` (used in cli.py) that loads them.

3"""

4 import typing

5 from dataclasses import dataclass

6 from importlib.metadata import EntryPoint, entry_points

7

8 from rich import print

9 from typer import Typer

10

11 from .core import (

12 AbstractConfig,

13 ApplicationState,

14 T_Command,

15 print_json,

16 run_tool,

17 state,

18 with_exit_code,

19)

20

21 __all__ = ["register", "run_tool", "PluginConfig", "print", "print_json"]

22

23

24 class PluginConfig(AbstractConfig):

25 """

26 Can be inherited in plugin to load in plugin-specific config.

27

28 The config class is a Singleton, which means multiple instances can be created, but they always have the same state.

29

30 Example:

31 @register()

32 class DemoConfig(PluginConfig):

33 required_arg: str

34 boolean_arg: bool

35 # ...

36

37

38 config = DemoConfig()

39

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 """

45

46 extras: dict[str, typing.Any]

47 state: typing.Optional[ApplicationState] # only with @register(with_state=True) or after self.attach_state

48

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 = {}

56

57 def attach_extra(self, name: str, obj: typing.Any) -> None:

58 """

59 Add a non-annotated variable.

60 """

61 self.extras[name] = obj

62

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)

69

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

74

75 def _get(self, key: str, strict: bool = True) -> typing.Any:

76 notfound = object()

77

78 value = getattr(self, key, notfound)

79 if value is not notfound:

80 return value

81

82 if self.extras:

83 value = self.extras.get(key, notfound)

84 if value is not notfound:

85 return value

86

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)

92

93 def _values(self) -> typing.Generator[typing.Any, typing.Any, None]:

94 yield from (self._get(k, False) for k in self._fields())

95

96 def __repr__(self) -> str:

97 """

98 Create a readable representation of this class with its data.

99

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})"

107

108 def __str__(self) -> str:

109 """

110 Alias for repr.

111 """

112 return repr(self)

113

114

115 T_PluginConfig = typing.Type[PluginConfig]

116

117 U_Wrappable = typing.Union[T_PluginConfig, T_Command]

118 T_Wrappable = typing.TypeVar("T_Wrappable", T_PluginConfig, T_Command)

119

120

121@dataclass()

122 class Registration(typing.Generic[T_Wrappable]):

123 wrapped: T_Wrappable

124

125 # Command:

126 add_to_all: bool

127 add_to_fix: bool

128

129 # Config:

130 with_state: bool

131 strict: bool

132 config_key: typing.Optional[str]

133

134 args: tuple[typing.Any, ...]

135 kwargs: dict[str, typing.Any]

136

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"

143

144

145 AnyRegistration = Registration[T_PluginConfig] | Registration[T_Command]

146

147# WeakValueDictionary() does not work since it removes the references too soon :(

148 registrations: dict[int, AnyRegistration] = {}

149

150

151 def _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 )

174

175 registrations[id(wrapped)] = registration

176 state.register_plugin(wrapped.__name__, registration)

177

178 return wrapped

179

180

181@typing.overload

182 def 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.

185

186 @register

187 def func(): ...

188

189 -> register(func) is called

190 """

191

192

193@typing.overload

194 def 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.

199

200 @register()

201 def func(): ...

202

203 -> register() is called

204 """

205

206

207 def 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.

221

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 ...

227

228 @register() # () are optional, but extra keyword arguments can be passed to configure the config.

229 class MyConfig(PluginConfig):

230 property: str

231 """

232

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)

235

236 if wrappable:

237 return inner(wrappable)

238 else:

239 return inner

240

241

242 T = typing.TypeVar("T")

243

244

245 class BoundMethodOf(typing.Protocol[T]):

246 """

247 Protocol to define properties that a bound method has.

248 """

249

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

253

254 def __call__(self, a: int) -> str: # pragma: no cover

255 """

256 Indicates this Protocol type can be called.

257 """

258

259

260 Unbound = typing.Callable[..., typing.Any]

261

262

263 def 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)

268

269

270@dataclass()

271 class PluginLoader:

272 app: Typer

273 with_exit_code: bool

274

275 def main(self) -> None:

276 """

277 Using importlib.metadata, discover available su6 plugins.

278

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)

288

289 self._cleanup()

290

291 def _cleanup(self) -> None:

292 # registrations.clear()

293 ...

294

295 def _load_plugin(self, plugin: EntryPoint) -> list[str]:

296 """

297 Look for typer instances and registered commands and configs in an Entrypoint.

298

299 [project.entry-points."su6"]

300 demo = "su6_plugin_demo.cli"

301

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()

307

308 for item in dir(plugin_module):

309 if item.startswith("_"):

310 continue

311

312 possible_command = getattr(plugin_module, item)

313

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)))

316

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

324

325 return result

326

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 {_}"]

336

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.

340

341 Example:

342 # pyproject.toml

343 [tool.su6.demo]

344 boolean-arg = true

345 optional-with-default = "overridden"

346

347 [tool.su6.demo.extra]

348 more = true

349 """

350 key = registration.config_key or name

351

352 cls = registration.wrapped

353 inst = cls()

354

355 if registration.with_state:

356 inst.attach_state(state)

357

358 if registration.strict is False:

359 inst._strict = False

360

361 state.attach_plugin_config(key, inst)

362 return [f"config {name}"]

363

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}"]

367

368

369 def 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.

372

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()

380

381

382# todo:

383# - add to 'all'

384# - add to 'fix'