Enforce ptype rules in Plane.multiply

andykee · Oct 19, 2023 · c8c3684 · c8c3684
1 parent b350657
commit c8c3684
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/lentil/plane.py b/lentil/plane.py
@@ -365,6 +365,11 @@ def multiply(self, wavefront, inplace=False):
             Updated wavefront
 
         """
+        if not _can_multiply_ptype(self.ptype, wavefront.ptype):
+            raise TypeError(f"can't multiply Wavefront with ptype " \
+                            f"'{wavefront.ptype}' by Plane with ptype " \
+                            f"'{self.ptype}'")
+
         pixelscale = _multiply_pixelscale(self.pixelscale, wavefront.pixelscale)
         shape = wavefront.shape if self.shape == () else self.shape
         data = wavefront.data

diff --git a/lentil/wavefront.py b/lentil/wavefront.py
@@ -75,7 +75,7 @@ def ptype(self):
     @ptype.setter
     def ptype(self, value):
         if lentil.ptype(value) == lentil.transform:
-            raise TypeError("invalid type ptype('transform') for Wavefront")
+            raise TypeError("invalid ptype 'transform' for Wavefront")
         else:
             self._ptype = lentil.ptype(value)