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 example 

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

85 

86 def __init__(self) -> None: 

87 self._full_values: list[tuple[float, PLUGIN_TYPE]] = [] 

88 self._partial_values: list[tuple[float, PLUGIN_TYPE]] = [] 

89 self._channel_values: list[tuple[float, PLUGIN_TYPE]] = [] 

90 

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

92 yield from (func for _, func in self._channel_values) 

93 

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

95 yield from (func for _, func in self._partial_values) 

96 

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

98 yield from (func for _, func in self._full_values) 

99 

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

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

102 pretty picture. 

103 

104 parameters 

105 ---------- 

106 order : `float` 

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

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

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

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

111 kind : `PluginType` 

112 This specifies what data the registered plugin expects 

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

114 

115 """ 

116 

117 def wrapper( 

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

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

120 match kind: 

121 case PluginType.PARTIAL: 

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

123 case PluginType.FULL: 

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

125 case PluginType.CHANNEL: 

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

127 return func 

128 

129 return wrapper 

130 

131 

132plugins = PluginsRegistry() 

133""" 

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

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

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

137needing any other import time logic elsewhere. 

138"""