ruff: avoid quadratic list summation

AmpelAstro · Nov 27, 2024 · 616a02a · 616a02a
1 parent 19a6947
commit 616a02a
Showing 1 changed file with 5 additions and 11 deletions.
diff --git a/ampel/ztf/view/ZTFT2Tabulator.py b/ampel/ztf/view/ZTFT2Tabulator.py
@@ -15,6 +15,7 @@
 from ampel.abstract.AbsT2Tabulator import AbsT2Tabulator
 from ampel.content.DataPoint import DataPoint
 from ampel.types import StockId
+from ampel.util.collections import ampel_iter
 from ampel.ztf.util.ZTFIdMapper import ZTFIdMapper
 
 ZTF_BANDPASSES = {
@@ -89,17 +90,10 @@ def get_jd(
 
     def get_stock_id(self, dps: Iterable[DataPoint]) -> set[StockId]:
         return set(
-            sum(
-                [
-                    list(stockid)
-                    if isinstance(stockid := el["stock"], Sequence)
-                    and not isinstance(stockid, str | bytes)
-                    else [stockid]
-                    for el in dps
-                    if "ZTF" in el["tag"]
-                ],
-                [],
-            )
+            stockid
+            for el in dps
+            if "ZTF" in el["tag"]
+            for stockid in ampel_iter(el["stock"])
         )
 
     def get_stock_name(self, dps: Iterable[DataPoint]) -> list[str]: