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

53 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-05 12:23 +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, field 

6from importlib.metadata import entry_points 

7 

8from typer import Typer 

9 

10from .core import T_Command, T_Command_Return, with_exit_code 

11 

12 

13@dataclass 

14class PluginRegistration: 

15 """ 

16 When using the @register decorator, a Registration is created. 

17 

18 `discover_plugins` will use this class to detect Registrations in a plugin module 

19 and `include_plugins` will add them to the top-level Typer app. 

20 """ 

21 

22 func: T_Command 

23 args: tuple[typing.Any, ...] = field(default_factory=tuple) 

24 kwargs: dict[str, typing.Any] = field(default_factory=dict) 

25 

26 @property 

27 def IS_SU6_REGISTRATION(self) -> bool: 

28 """ 

29 Used to detect if a variable is a Registration. 

30 

31 Even when isinstance() does not work because it's stored in memory in two different locations 

32 (-> e.g. in a plugin). 

33 

34 See also, is_registration 

35 """ 

36 return True 

37 

38 def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> T_Command_Return: 

39 """ 

40 You can still use a Plugin Registration as a normal function. 

41 """ 

42 return self.func(*args, **kwargs) 

43 

44 

45T_Inner = typing.Callable[[T_Command], PluginRegistration] 

46 

47 

48@typing.overload 

49def register( 

50 func_outer: None = None, 

51 *a_outer: typing.Any, 

52 **kw_outer: typing.Any, 

53) -> T_Inner: 

54 """ 

55 If func outer is None, a callback will be created that will return a Registration later. 

56 

57 Example: 

58 @register() 

59 def command 

60 """ 

61 

62 

63@typing.overload 

64def register( 

65 func_outer: T_Command, 

66 *a_outer: typing.Any, 

67 **kw_outer: typing.Any, 

68) -> PluginRegistration: 

69 """ 

70 If func outer is a command, a registration will be created. 

71 

72 Example: 

73 @register 

74 def command 

75 """ 

76 

77 

78def register( 

79 func_outer: T_Command | None = None, 

80 *a_outer: typing.Any, 

81 **kw_outer: typing.Any, 

82) -> PluginRegistration | T_Inner: 

83 """ 

84 Decorator used to add a top-level command to `su6`. 

85 

86 Can either be used as @register() or @register. 

87 

88 Args: 

89 func_outer: wrapped method 

90 a_outer: arguments passed to @register(arg1, arg2, ...) - should probably not be used! 

91 kw_outer: keyword arguments passed to @register(name=...) - will be passed to Typer's @app.command 

92 """ 

93 if func_outer: 

94 # @register 

95 # def func 

96 return PluginRegistration(func_outer, a_outer, kw_outer) 

97 

98 # @functools.wraps(func_outer) 

99 def inner(func_inner: T_Command) -> PluginRegistration: 

100 # @register() 

101 # def func 

102 

103 # combine args/kwargs from inner and outer, just to be sure they are passed. 

104 return PluginRegistration(func_inner, a_outer, kw_outer) 

105 

106 return inner 

107 

108 

109# list of registrations 

110T_Commands = list[PluginRegistration] 

111 

112# key: namespace 

113# value: app instance, docstring for 'help' 

114T_Namespaces = dict[str, tuple[Typer, str]] 

115 

116 

117def is_registration(something: typing.Any) -> bool: 

118 """ 

119 Pytest might freak out if some package is pip installed and Registration exists locally. 

120 

121 This method uses IS_SU6_REGISTRATION to check if the types actually match. 

122 """ 

123 return getattr(something, "IS_SU6_REGISTRATION", False) 

124 

125 

126def discover_plugins() -> tuple[T_Namespaces, T_Commands]: 

127 """ 

128 Using importlib.metadata, discover available su6 plugins. 

129 

130 Example: 

131 # pyproject.toml 

132 # https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata 

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

134 demo = "su6_plugin_demo.cli" # <- CHANGE ME 

135 """ 

136 discovered_namespaces = {} 

137 discovered_commands = [] 

138 discovered_plugins = entry_points(group="su6") 

139 for plugin in discovered_plugins: 

140 plugin_module = plugin.load() 

141 

142 for item in dir(plugin_module): 

143 if item.startswith("_"): 

144 continue 

145 

146 possible_command = getattr(plugin_module, item) 

147 

148 if isinstance(possible_command, Typer): 

149 discovered_namespaces[plugin.name] = (possible_command, plugin_module.__doc__) 

150 elif is_registration(possible_command): 

151 discovered_commands.append(possible_command) 

152 

153 return discovered_namespaces, discovered_commands 

154 

155 

156def include_plugins(app: Typer, _with_exit_code: bool = True) -> None: 

157 """ 

158 Discover plugins using discover_plugins and add them to either global namespace or as a subcommand. 

159 

160 Args: 

161 app: the top-level Typer app to append commands to 

162 _with_exit_code: should the @with_exit_code decorator be applied to the return value of the command? 

163 """ 

164 namespaces, commands = discover_plugins() 

165 

166 for namespace, (subapp, doc) in namespaces.items(): 

167 # adding subcommand 

168 app.add_typer(subapp, name=namespace, help=doc) 

169 

170 for registration in commands: 

171 if _with_exit_code: 

172 registration.func = with_exit_code()(registration.func) 

173 

174 # adding top-level commands 

175 app.command(*registration.args, **registration.kwargs)(registration.func) 

176 

177 

178# todo: 

179# - add to 'all' 

180# - add to 'fix' 

181# - adding config keys