Coverage for python / lsst / pipe / tasks / prettyPictureMaker / _plugins.py: 56%

1# This file is part of pipe_tasks. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ("PluginsRegistry", "plugins") 

25 

26from enum import Enum, auto 

27from collections.abc import Callable 

28from typing import TYPE_CHECKING, Generator 

29from lsst.pipe.base import PipelineTaskConfig 

30 

31 

32if TYPE_CHECKING: 

33 from numpy.typing import NDArray 

34 from collections.abc import Mapping 

35 

36 PLUGIN_TYPE = Callable[[NDArray, NDArray, Mapping[str, int]], NDArray] 

37 

38 

39class PluginType(Enum): 

40 """Enumeration to mark the type of data a plugin expects to work on""" 

41 

42 CHANNEL = auto() 

43 """A plugin of this type expects to work on an individual channel 

44 from a partial region of a mosaic, such as a `patch`. 

45 """ 

46 PARTIAL = auto() 

47 """A pluging that expects to work on a 3 channel image that is 

48 a partial region of a mosaic, such as a `patch`. 

49 """ 

50 FULL = auto() 

51 """FULL plugins operate on a 3 channel image corresponding to 

52 a complete mosaic. 

53 """ 

54 

55 

56class PluginsRegistry: 

57 """A class to serve as a registry for all pretty picture manipulation 

58 plugins. 

59 

60 This class should not be instantiated directly other than the one 

61 instantiation in this module. 

62 

63 Examples 

64 -------- 

65 Using this registry to create a plugin would look somehting like the 

66 following. 

67 

68 >>> @plugins.register(1, PluginType.PARTIAL) 

69 >>> def fixNoData( 

70 >>> image: NDArray, 

71 >>> mask: NDArray, 

72 >>> maskDict: Mapping[str, int] 

73 >>> ) -> NDArray: 

74 >>> m = (mask & 2 ** maskDict["NO_DATA"]).astype(bool) 

75 >>> for i in range(3): 

76 >>> image[:, :, i] = cv2.inpaint( 

77 >>> image[:, :, i].astype(np.float32), 

78 >>> m.astype(np.uint8), 

79 >>> 3, 

80 >>> cv2.INPAINT_TELEA 

81 >>> ).astype(image.dtype) 

82 >>> return image 

83 

84 Parameters 

85 ---------- 

86 None 

87 

88 """ 

89 

90 def __init__(self) -> None: 

91 self._full_values: list[tuple[float, Callable]] = [] 

92 self._partial_values: list[tuple[float, Callable]] = [] 

93 self._channel_values: list[tuple[float, Callable]] = [] 

94 

95 def channel(self) -> Generator[PLUGIN_TYPE, None, None]: 

96 """Yield generators of channel plugins. 

97 

98 Returns 

99 ------- 

100 gen : `~collections.abc.Iterator` 

101 Generator of channel plugins. 

102 """ 

103 return (func for _, func in self._channel_values) 

104 

105 def partial(self) -> Generator[PLUGIN_TYPE, None, None]: 

106 """Yield generators of partial plugins. 

107 

108 Returns 

109 ------- 

110 gen : `~collections.abc.Iterator` 

111 Generator of partial plugins. 

112 """ 

113 return (func for _, func in self._partial_values) 

114 

115 def full(self) -> Generator[PLUGIN_TYPE, None, None]: 

116 """Yield generators of full plugins. 

117 

118 Returns 

119 ------- 

120 gen : `~collections.abc.Iterator` 

121 Generator of full plugins. 

122 """ 

123 return (func for _, func in self._full_values) 

124 

125 def register(self, order: float, kind: PluginType) -> Callable: 

126 """Register a plugin which is to be run when producing a 

127 pretty picture. 

128 

129 Parameters 

130 ---------- 

131 order : `float` 

132 This determines in what order plugins will be run. For 

133 example, if plugin A specifies order 2, and plugin B 

134 specifies order 1, and both are the same ``kind`` of 

135 plugin type, plugin B will be run before plugin A. 

136 kind : `PluginType` 

137 This specifies what data the registered plugin expects 

138 to run on, a channel, a partial image, or a full mosaic. 

139 

140 Returns 

141 ------- 

142 wrapper : `Callable` 

143 Decorator function for registering the plugin. 

144 """ 

145 

146 def wrapper( 

147 func: Callable[[NDArray, NDArray, Mapping[str, int], PipelineTaskConfig], NDArray], 

148 ) -> Callable[[NDArray, NDArray, Mapping[str, int], PipelineTaskConfig], NDArray]: 

149 """Wrapper decorator for registering plugin functions. 

150 

151 Parameters 

152 ---------- 

153 func : `Callable` 

154 Plugin function being registered. 

155 

156 Returns 

157 ------- 

158 func : `Callable` 

159 The same plugin function, now registered in the registry. 

160 """ 

161 match kind: 

162 case PluginType.PARTIAL: 

163 self._partial_values.append((order, func)) 

164 case PluginType.FULL: 

165 self._full_values.append((order, func)) 

166 case PluginType.CHANNEL: 

167 self._channel_values.append((order, func)) 

168 return func 

169 

170 return wrapper 

171 

172 

173plugins = PluginsRegistry() 

174""" 

175This is the only instance of the plugin registry there should be. Users 

176should import from here and use the register method as a decorator to 

177register any plugins. Or, preferably, add them to this file to avoid 

178needing any other import time logic elsewhere. 

179"""