Parametric type hinting for control paths

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

Start by inspecting the definitions of AbstractTerm, _ControlTerm, and WeaklyDiagonalControlTerm, then review how their control annotations are currently expressed. Check whether the proposed generic type parameters can distinguish VirtualBrownianPath from UnsafeBrownianPath without breaking existing uses. Done means the supported typing behavior and its scope are documented or covered by appropriate checks.