Fh/weeeee

scopesim/effects/apertures.py

@@ -11,7 +11,7 @@
 from astropy import units as u
 from astropy.table import Table
 
-from .effects import Effect
+from .effects import Effect, WheelEffect
 from ..optics import image_plane_utils as imp_utils
 from ..base_classes import FOVSetupBase
 
@@ -384,7 +384,7 @@
     #     return self.get_apertures(item)[0]
 
 
-class SlitWheel(Effect):
+class SlitWheel(WheelEffect, Effect):
     """
     Selection of predefined spectroscopic slits and possibly other field masks.
 
@@ -426,35 +426,35 @@
 
     """
 
-    required_keys = {"slit_names", "filename_format", "current_slit"}
     z_order: ClassVar[tuple[int, ...]] = (80, 280, 580)
     report_plot_include: ClassVar[bool] = False
     report_table_include: ClassVar[bool] = True
     report_table_rounding: ClassVar[int] = 4
-    _current_str = "current_slit"
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        check_keys(kwargs, self.required_keys, action="error")
+    _item_cls: ClassVar[type] = ApertureMask
+    _item_str: ClassVar[str] = "slit"
 
-        params = {
-            "path": "",
-        }
-        self.meta.update(params)
-        self.meta.update(kwargs)
+    def __init__(self, slit_names, filename_format, current_slit, **kwargs):
+        super().__init__(kwargs=kwargs)
+
+        self.meta["path"] = ""
+        self.meta["filename_format"] = filename_format
 
         path = self._get_path()
-        self.slits = {}
-        for name in from_currsys(self.meta["slit_names"], self.cmds):
-            kwargs["name"] = name
+        for name in from_currsys(slit_names, self.cmds):
             fname = str(path).format(name)
-            self.slits[name] = ApertureMask(filename=fname, **kwargs)
+            self.items[name] = self._item_cls(filename=fname, name=name)
 
+        self.current_item_name = from_currsys(current_slit, self.cmds)
         self.table = self.get_table()
 
-    def apply_to(self, obj, **kwargs):
-        """Use apply_to of current_slit."""
-        return self.current_slit.apply_to(obj, **kwargs)
+    @property
+    def slits(self):
+        return self.items
+
+    @slits.setter
+    def slits(self, value):
+        self.items = value
 
     def fov_grid(self, which="edges", **kwargs):
         """See parent docstring."""
@@ -464,37 +464,27 @@
 
     def change_slit(self, slitname=None):
         """Change the current slit."""
-        if not slitname or slitname in self.slits.keys():
-            self.meta["current_slit"] = slitname
-            self.include = slitname
-        else:
-            raise ValueError("Unknown slit requested: " + slitname)
+        self.change_item(slitname)
 
     def add_slit(self, newslit, name=None):
         """
-        Add a slit to the SlitWheel.
+        Add a Slit to the SlitWheel.
 
         Parameters
         ----------
-        newslit : Slit
-        name : string
-           Name to be used for the new slit. If ``None``, a name from
-           the newslit object is used.
+        new_item : Slit
+            Slit instance to be added.
+        item_name : str, optional
+            Name to be used for the new Slit. If `None`, the name is taken from
+            the Slits's name. The default is None.
+
         """
-        if name is None:
-            name = newslit.display_name
-        self.slits[name] = newslit
+        self.add_item(newslit, name)
 
     @property
     def current_slit(self):
         """Return the currently used slit."""
-        currslit = from_currsys(self.meta["current_slit"], self.cmds)
-        if not currslit:
-            return False
-        return self.slits[currslit]
-
-    def __getattr__(self, item):
-        return getattr(self.current_slit, item)
+        return self.current_item
 
     def get_table(self):
         """
@@ -503,8 +493,8 @@
         Width is defined as the extension in the y-direction, length in the
         x-direction. All values are in milliarcsec.
         """
-        names = list(self.slits.keys())
-        slits = self.slits.values()
+        names = list(self.items.keys())
+        slits = self.items.values()
         xmax = np.array([slit.data["x"].max() * u.Unit(slit.meta["x_unit"])
                          .to(u.mas) for slit in slits])
         xmin = np.array([slit.data["x"].min() * u.Unit(slit.meta["x_unit"])

scopesim/effects/effects.py

@@ -2,16 +2,19 @@
 """Contains base class for effects."""
 
 from pathlib import Path
-from collections.abc import Mapping, MutableMapping
+from collections.abc import Mapping, MutableMapping, Iterable
 from dataclasses import dataclass, field, InitVar, fields
-from typing import NewType, ClassVar
+from typing import NewType, ClassVar, Any
 
 from .data_container import DataContainer
 from .. import base_classes as bc
-from ..utils import from_currsys, write_report
+from ..utils import from_currsys, write_report, get_logger
 from ..reports.rst_utils import table_to_rst
 
 
+logger = get_logger(__name__)
+
+
 # FIXME: This docstring is out-of-date for several reasons:
 #   - Effects can act on objects other than Source (eg FOV, IMP, DET)
 #   - fov_grid is outdated
@@ -141,11 +144,7 @@
 
     @property
     def display_name(self) -> str:
-        name = self.meta.get("name", self.meta.get("filename", "<untitled>"))
-        if not hasattr(self, "_current_str"):
-            return name
-        current_str = from_currsys(self.meta[self._current_str], self.cmds)
-        return f"{name} : [{current_str}]"
+        return self.meta.get("name", self.meta.get("filename", "<untitled>"))
 
     @property
     def meta_string(self) -> str:
@@ -351,7 +350,13 @@
                     else:
                         value = from_currsys(self.meta, self.cmds)
                 else:
-                    value = self.meta[item.removeprefix("#")]
+                    try:
+                        value = self.meta[item.removeprefix("#")]
+                    except KeyError as err:
+                        try:
+                            value = getattr(self, item.removeprefix("#"))
+                        except AttributeError:
+                            raise err
             else:
                 value = self.meta
         else:
@@ -364,3 +369,82 @@
             return None
         return Path(self.meta["path"],
                     from_currsys(self.meta["filename_format"], self.cmds))
+
+
+# TODO: Maybe make this a MutableMapping instead and have items as the map?
+@dataclass(kw_only=True, eq=False)
+class WheelEffect:
+    """Base class for wheel-type effects."""
+
+    items: dict[str, Effect] = field(
+        default_factory=dict, init=False, repr=False
+    )
+    current_item_name: str | None = None
+    kwargs: InitVar[Any] = None  # HACK: remove this once proper dataclass effs
+
+    _item_cls: ClassVar[type] = Effect
+    _item_str: ClassVar[str] = "item"
+
+    def __post_init__(self, kwargs):
+        # TODO: remove this once Effect is a proper dataclass
+        if kwargs is None:
+            kwargs = {}
+        super().__init__(**kwargs)
+
+    @property
+    def display_name(self) -> str:
+        return f"{super().display_name} : [{self.current_item_name}]"
+
+    def apply_to(self, obj, **kwargs):
+        """Use apply_to of current item."""
+        return self.current_item.apply_to(obj, **kwargs)
+
+    def change_item(self, item_name) -> None:
+        """Change the current item."""
+        if item_name not in self.items:
+            # current=False is sometimes used to disable effect
+            # TODO: do we really want this?
+            if item_name is False:  # need explicit check here to avoid ""
+                self.include = False
+                return
+            raise ValueError(
+                f"Unknown {self._item_str} requested: {item_name}"
+            )
+        self.current_item_name = item_name
+
+    def add_item(self, new_item, item_name: str | None = None) -> None:
+        """
+        Add an item to the Wheel.
+
+        Parameters
+        ----------
+        new_item : Effect
+            Effect subclass item to be added.
+        item_name : str, optional
+            Name to be used for the new item. If `None`, the name is taken from
+            the item's name. The default is None.
+
+        """
+        if (name := item_name or new_item.display_name) in self.items:
+            logger.warning("%s already in wheel, overwriting", name)
+        self.items[name] = new_item
+
+    @property
+    def current_item(self):
+        """Return the currently selected item (`None` if not set)."""
+        # TODO: do we really want this?
+        if not self.include:
+            return False
+        return self.items.get(self.current_item_name, None)
+
+    @current_item.setter
+    def current_item(self):
+        raise AttributeError(
+            f"{self.__class__.__name__}.current_{self._item_str} cannot be "
+            f"set directly. Use {self.__class__.__name__}.change_"
+            f"{self._item_str}({self._item_str}_name) instead."
+        )
+
+    def __getattr__(self, key):
+        # TODO: reevaluate the need for this...
+        return getattr(self.current_item, key)

scopesim/effects/spectral_trace_list.py

@@ -14,7 +14,7 @@
 from astropy.io import fits
 from astropy.table import Table
 
-from .effects import Effect
+from .effects import Effect, WheelEffect
 from .ter_curves import FilterCurve
 from .spectral_trace_list_utils import SpectralTrace, make_image_interpolations
 from ..optics.image_plane_utils import header_from_list_of_xy
@@ -408,7 +408,7 @@
         self.spectral_traces[key] = value
 
 
-class SpectralTraceListWheel(Effect):
+class SpectralTraceListWheel(WheelEffect, Effect):
     """
     A Wheel-Effect object for selecting between multiple gratings/grisms.
 
@@ -466,46 +466,38 @@
 
     """
 
-    required_keys = {
-        "trace_list_names",
-        "filename_format",
-        "current_trace_list",
-    }
     z_order: ClassVar[tuple[int, ...]] = (70, 270, 670)
     report_plot_include: ClassVar[bool] = True
     report_table_include: ClassVar[bool] = True
     report_table_rounding: ClassVar[int] = 4
-    _current_str = "current_trace_list"
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        check_keys(kwargs, self.required_keys, action="error")
+    _item_cls: ClassVar[type] = SpectralTraceList
+    _item_str: ClassVar[str] = "trace_list"
 
-        params = {
-            "path": "",
-        }
-        self.meta.update(params)
-        self.meta.update(kwargs)
+    def __init__(self, trace_list_names, filename_format, current_trace_list,
+                 **kwargs):
+        super().__init__(kwargs=kwargs)
+
+        self.meta["path"] = ""
+        self.meta["filename_format"] = filename_format
 
         path = self._get_path()
-        self.trace_lists = {}
-        if "name" in kwargs:
-            kwargs.pop("name")
-        for name in from_currsys(self.meta["trace_list_names"], self.cmds):
+        kwargs.pop("name", None)
+        for name in from_currsys(trace_list_names, self.cmds):
             fname = str(path).format(name)
-            self.trace_lists[name] = SpectralTraceList(filename=fname,
-                                                       name=name,
-                                                       **kwargs)
+            self.items[name] = self._item_cls(filename=fname, name=name,
+                                              **kwargs)
 
-    def apply_to(self, obj, **kwargs):
-        """Use apply_to of current trace list."""
-        return self.current_trace_list.apply_to(obj, **kwargs)
+        self.current_item_name = from_currsys(current_trace_list, self.cmds)
+
+    @property
+    def trace_lists(self):
+        return self.items
+
+    @trace_lists.setter
+    def trace_lists(self, value):
+        self.items = value
 
     @property
     def current_trace_list(self):
-        trace_list_eff = None
-        trace_list_name = from_currsys(self.meta["current_trace_list"],
-                                       self.cmds)
-        if trace_list_name is not None:
-            trace_list_eff = self.trace_lists[trace_list_name]
-        return trace_list_eff
+        return self.current_item