Fixing mypy errors in optika.sags

sun-data · Aug 14, 2023 · f1b8213 · f1b8213
1 parent 1a0fee5
commit f1b8213
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/optika/sags.py b/optika/sags.py
@@ -13,9 +13,9 @@
     "ToroidalSag",
 ]
 
-RadiusT = TypeVar("RadiusT", bound=na.ScalarLike)
-ConicT = TypeVar("ConicT", bound=na.ScalarLike)
-RadiusOfRotationT = TypeVar("RadiusOfRotationT", bound=na.ScalarLike)
+RadiusT = TypeVar("RadiusT", bound=float | np.ndarray | na.AbstractScalar)
+ConicT = TypeVar("ConicT", bound=float | np.ndarray | na.AbstractScalar)
+RadiusOfRotationT = TypeVar("RadiusOfRotationT", bound=float | np.ndarray | na.AbstractScalar)
 
 
 @dataclasses.dataclass
@@ -85,7 +85,7 @@ def __call__(
     def normal(
         self,
         position: na.AbstractCartesian2dVectorArray,
-    ) -> na.Cartesian3dVectorArray:
+    ) -> na.AbstractCartesian3dVectorArray:
         radius = self.radius
         c = self.curvature
 
@@ -144,7 +144,7 @@ def __call__(
     def normal(
         self,
         position: na.AbstractCartesian2dVectorArray,
-    ) -> na.Cartesian3dVectorArray:
+    ) -> na.AbstractCartesian3dVectorArray:
         radius = self.radius
         c = self.curvature
         conic = self.conic
@@ -205,7 +205,7 @@ def __call__(
     def normal(
         self,
         position: na.AbstractCartesian2dVectorArray,
-    ) -> na.AbstractScalar:
+    ) -> na.AbstractCartesian3dVectorArray:
         c = self.curvature
         r = self.radius_of_rotation