Parametric type hinting for control paths

[tttc3]

I would like to type hint that my custom term expects a control which has a certain type of path. E.G. only VirtualBrownianTree not UnsafeBrownianPath.

Ideally the code would look something like this:

from diffrax import AbstractTerm, ODETerm, WeaklyDiagonalControlTerm

class CustomTerm(AbstractTerm)
    ode: ODETerm
    cde: WeaklyDiagonalControlTerm[VirtualBrownianPath]

I'm not sure if this is already possible, or would require a change to the code. My thought was to do something like the following:

class _ControlTerm[T: Union[AbstractPath, Callable]](AbstractTerm):
    vector_field: ...
    control: T = ...

class WeaklyDiagonalControlTerm[T](_ControlTerm[T]):
    ...

I'm not very experienced with the more advanced aspects of type hinting so perhaps there is a better way, or this is already possible? Any thought are much appreciated

[patrick-kidger]

This seems reasonable to me! I'd be happy to take a PR adding this. Note that we don't yet insist on Python 3.12 and above, so I think we'd annotate this as something like

_Control = TypeVar("_Control", bound=AbstractPath)

class _ControlTerm(AbstractTerm, Generic[_Control]):
    control: _Control = ...

    def __init__(self, ..., control: Union[AbstractPath, Callable]):
        self.control = _callable_to_path(control)