From 3d5040221c0249cd343e8f0e0cecbb8dcb0cf870 Mon Sep 17 00:00:00 2001
From: Dan Weatherill <daniel.weatherill@physics.ox.ac.uk>
Date: Thu, 13 Jun 2024 13:36:22 +0100
Subject: [PATCH] refactor into base class for Repack actions

---
 .../tools/atools/calibQuantityProfile.py      | 51 +++++++++----------
 1 file changed, 25 insertions(+), 26 deletions(-)

diff --git a/python/lsst/analysis/tools/atools/calibQuantityProfile.py b/python/lsst/analysis/tools/atools/calibQuantityProfile.py
index 545901d65..5340afa8a 100644
--- a/python/lsst/analysis/tools/atools/calibQuantityProfile.py
+++ b/python/lsst/analysis/tools/atools/calibQuantityProfile.py
@@ -31,7 +31,7 @@
     "CalibPtcCovarScatterTool",
 )
 
-from typing import cast, TypeVar, Callable, Any
+from typing import cast, TypeVar, Callable, Any, TypeAlias
 
 from lsst.pex.config import Field, ListField
 from lsst.pex.config.configurableActions import ConfigurableActionField
@@ -59,22 +59,23 @@
     15: "C17",
 }
 
+#Dummy class just so we can get the type annotations correct
+RepackerLoopFun: TypeAlias = Callable[None, [KeyedData, Any, Any, Any, Any]]
 
-RepackerT = TypeVar("RepackerT", bound=[PrepRepacker, SingleValueRepacker])
-RepackerLoopFun = Callable[None, [KeyedData, Any, Any, Any, Any]]
-def _repack_loop_helper(obj: RepackerT, repackfun: RepackerLoopFun, data: KeyedData) -> KeyedData:
-    repackedData: dict[str, Vector] = {}
-    quantitiesData = [data[_] for _ in obj.quantityKey]
-
-    for p, d in zip(data[obj.panelKey], data[obj.dataKey]):
-        for qName, qD in zip(obj.quantityKey, quantitiesData):
-            RepackerLoopFun(repackedData, p, d, qName, qD)
-    return repackedData
+class BaseRepacker(KeyedDataAction):
+    """Base class for Data Repacking actions. Essentially Just adds some helper functions"""
+    def _repack_loop_helper(self, repackfun: RepackerLoopFun, data: KeyedData) -> KeyedData:
+        repackedData: dict[str, Vector] = {}
+        quantitiesData = [data[_] for _ in obj.quantityKey]
 
+        for p, d in zip(data[obj.panelKey], data[obj.dataKey]):
+            for qName, qD in zip(obj.quantityKey, quantitiesData):
+                RepackerLoopFun(repackedData, p, d, qName, qD)
+        return repackedData
 
 
 
-class PrepRepacker(KeyedDataAction):
+class PrepRepacker(BaseRepacker):
     """Prep action to repack data."""
 
     panelKey = Field[str](
@@ -91,7 +92,7 @@ def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
         def rp_loop(rpData: KeyedData, panel, data, quantityName, quantityData):
             qName = f"{panel}_{data}_{quantityName}"
             rpData[qName] = quantityData
-        repackedData = _repack_loop_helper(self, rp_loop, data) 
+        repackedData = self._repack_loop_helper(self, rp_loop, data) 
         return repackedData
 
     def getInputSchema(self) -> KeyedDataSchema:
@@ -105,7 +106,7 @@ def addInputSchema(self, inputSchema: KeyedDataSchema) -> None:
         pass
 
 
-class SingleValueRepacker(KeyedDataAction):
+class SingleValueRepacker(BaseRepacker):
     """Prep action to repack data."""
 
     panelKey = Field[str](
@@ -119,22 +120,20 @@ class SingleValueRepacker(KeyedDataAction):
         minLength=1, optional=False)
 
     def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
-        uniquePanelKeys: list = list(set(data[self.panelKey]))
-
         repackedData: dict[str, Vector] = {}
-        uniquePanelKeys: list = list(set(data[self.panelKey]))
-
-        # Loop over data vector to repack information as it is expected.
-        for uk in uniquePanelKeys:
-            repackedData[f"{uk}_x"] = []
-            for q in self.quantityKey:
-                repackedData[f"{uk}_{q}"] = []
 
         def rp_loop(rpData: KeyedData, panel, data, quantityName, quantityData):
-            rpData[f"{panel}_x"].append(data)
-            rpData[f"{panel}_{quantityName}"] = quantityData
+            if (xlab := f"{panel}_x") not in rpData:
+                rpData[xlab] = [data]
+            else:
+                rpData[xlab].append(data)
+
+            if (lab := f"{panel}_{quantityName}") not in rpData:
+                rpData[lab] = [quantityData]
+            else:
+                rpData[lab].append(quantityData)
 
-        repackedData: KeyedData = _repack_loop_helper(self, rp_loop, data)
+        repackedData = self._repack_loop_helper(self, rp_loop, data)
         return repackedData
 
     def getInputSchema(self) -> KeyedDataSchema: