diff --git a/src/openflexure_microscope_server/utilities.py b/src/openflexure_microscope_server/utilities.py index 318e6457..b7377f9a 100644 --- a/src/openflexure_microscope_server/utilities.py +++ b/src/openflexure_microscope_server/utilities.py @@ -1,6 +1,15 @@ """Utility functions and classes.""" -from typing import TypeVar, Callable, ParamSpec, Optional, Any, Concatenate, Self +from typing import ( + TypeVar, + Callable, + ParamSpec, + Optional, + Any, + Concatenate, + Self, + overload, +) import os import re import sys @@ -11,6 +20,7 @@ import tomllib from functools import wraps from pydantic import BaseModel +import numpy as np T = TypeVar("T") P = ParamSpec("P") @@ -329,10 +339,20 @@ def _get_version_from_toml(toml_path: str) -> str: return "Undefined" -def quadratic(x: float, a: float, b: float, c: float) -> float: +# Let MyPy know that the output type matches x. +@overload +def quadratic(x: np.ndarray, a: float, b: float, c: float) -> np.ndarray: ... +@overload +def quadratic(x: float, a: float, b: float, c: float) -> float: ... + + +def quadratic( + x: float | np.ndarray, a: float, b: float, c: float +) -> float | np.ndarray: """Quadratic function. Used for predicting z. - :param x: The points at which to evaluate the quadratic. + :param x: The point or points at which to evaluate the quadratic. This can be a + float or a numpy array. The return will be the same type. :param a: The coefficient of x^2. :param b: The coefficient of x. :param c: The constant coefficient.