enable Quantity typing with syntax of Quantity[unit]

brainunit/_base.py

@@ -20,7 +20,7 @@
 from contextlib import contextmanager
 from copy import deepcopy
 from functools import wraps, partial
-from typing import Union, Optional, Sequence, Callable, Tuple, Any, List, Dict, cast
+from typing import Union, Optional, Sequence, Callable, Tuple, Any, List, Dict, cast, TypeVar, Generic
 
 import jax
 import jax.numpy as jnp
@@ -73,6 +73,8 @@
 PyTree = Any
 _all_slice = slice(None, None, None)
 compat_with_equinox = False
+A = TypeVar('A')
+
 
 
 def compatible_with_equinox(mode: bool = True):
@@ -2135,7 +2137,7 @@ def _element_not_quantity(x):
 
 
 @register_pytree_node_class
-class Quantity:
+class Quantity(Generic[A]):
     """
     The `Quantity` class represents a physical quantity with a mantissa and a unit.
     It is used to represent all physical quantities in ``BrainUnit``.

brainunit/_base_test.py

@@ -29,6 +29,7 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_equal
+from typing import Union
 
 import brainunit as u
 from brainunit._base import (
@@ -900,6 +901,16 @@ def test_to(self):
         print(x.to(u.volt))
         print(x.to(u.uvolt))
 
+    def test_quantity_type(self):
+        def f1(a: u.Quantity[u.ms]) -> u.Quantity[u.mV]:
+            return a
+
+        def f2(a: u.Quantity[Union[u.ms, u.mA]]) -> u.Quantity[u.mV]:
+            return a
+
+        def f3(a: u.Quantity[Union[u.ms, u.mA]]) -> u.Quantity[Union[u.mV, u.ms]]:
+            return a
+
 
 class TestNumPyFunctions(unittest.TestCase):
     def test_special_case_numpy_functions(self):
@@ -1468,12 +1479,6 @@ def test_pickle():
         print(b)
 
 
-
-
-
-
-
-
 def test_str_repr():
     """
     Test that str representations do not raise any errors and that repr