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

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 

7 

8from rich import print 

9from typer import Typer 

10 

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

24class 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 

115T_PluginConfig = typing.Type[PluginConfig] 

116 

117U_Wrappable = typing.Union[T_PluginConfig, T_Command] 

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

119 

120 

121@dataclass() 

122class 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 

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

146 

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

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

149 

150 

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 ) 

174 

175 registrations[id(wrapped)] = registration 

176 state.register_plugin(wrapped.__name__, registration) 

177 

178 return wrapped 

179 

180 

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. 

185 

186 @register 

187 def func(): ... 

188 

189 -> register(func) is called 

190 """ 

191 

192 

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. 

199 

200 @register() 

201 def func(): ... 

202 

203 -> register() is called 

204 """ 

205 

206 

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. 

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 

242T = typing.TypeVar("T") 

243 

244 

245class 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 

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

261 

262 

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) 

268 

269 

270@dataclass() 

271class 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 

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. 

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'