Normalize real roots returned by CurveFunction

sketchresponse/grader_lib/CurveFunction.py

@@ -171,11 +171,13 @@ def get_value_at(self, xval: float) -> float | None:
     def get_x_for_yval(self, yval: float) -> list[float]:
         p = np.poly1d(self.y)
         tvals = (p - yval).roots
-        xvals = [self.get_x_for_t(t) for t in tvals]
-        xvals = [val for val in xvals if np.imag(val) == 0]
-        return xvals
+        return [
+            float(np.real(self.get_x_for_t(t)))
+            for t in tvals
+            if abs(np.imag(t)) < self.tolerance["imag_threshold"]
+        ]
 
-    def get_x_for_t(self, t: float) -> float:
+    def get_x_for_t(self, t: complex) -> complex:
         x = (
             (1 - t) ** 3 * self.p0[0]
             + 3 * (1 - t) ** 2 * t * self.p1[0]

test_grader_lib/testSplineFunctionMethods.py

@@ -2,6 +2,7 @@
 from typing import cast
 
 from sketchresponse.grader_lib.Axis import Axis
+from sketchresponse.grader_lib.CurveFunction import CurveFunction
 from sketchresponse.grader_lib.SplineFunction import SplineFunction
 from sketchresponse.types import SketchGrader
 
@@ -29,5 +30,42 @@ def test_custom_tolerance_is_applied(self):
         self.assertEqual(len(sf.functions), 1)
 
 
+class TestCurveFunction(unittest.TestCase):
+    def test_get_x_for_yval_returns_real_floats(self):
+        xaxis = Axis([0, 1], 100)
+        yaxis = Axis([1, 0], 100)
+        path_info: list[list[float]] = [
+            [0, 100],
+            [100 / 3, 100],
+            [200 / 3, 100],
+            [100, 0],
+        ]
+        grader = cast(SketchGrader, {"tolerance": 10, "debug": False})
+        curve = CurveFunction(xaxis, yaxis, path_info, grader, "test-tool")
+
+        xvals = curve.get_x_for_yval(0.125)
+
+        self.assertEqual(len(xvals), 1)
+        self.assertAlmostEqual(xvals[0], 0.5)
+        self.assertListEqual([float] * len(xvals), [type(x) for x in xvals])
+
+    def test_get_x_for_yval_filters_complex_roots_before_scaling(self):
+        xaxis = Axis([0, 1e-6], 100)
+        yaxis = Axis([1, 0], 100)
+        path_info: list[list[float]] = [
+            [0, 100],
+            [100, 100],
+            [0, 100],
+            [100, 0],
+        ]
+        grader = cast(SketchGrader, {"tolerance": 10, "debug": False})
+        curve = CurveFunction(xaxis, yaxis, path_info, grader, "test-tool")
+
+        xvals = curve.get_x_for_yval(0.125)
+
+        self.assertEqual(len(xvals), 1)
+        self.assertAlmostEqual(xvals[0], 5e-7)
+
+
 if __name__ == "__main__":
     unittest.main()