fix typing errors

dendritex/_base.py

@@ -16,6 +16,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Optional, Dict, Sequence, Callable, NamedTuple, Tuple
 
 import brainstate as bst

dendritex/_base_test.py

@@ -15,7 +15,7 @@
 
 import unittest
 
+
 class TestIon(unittest.TestCase):
   def test1(self):
     pass
-

dendritex/_integrators.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 import brainstate as bst
 import brainunit as bu
 import jax
@@ -37,7 +39,6 @@ def tree_map(f, tree, *rest):
   return jax.tree.map(f, tree, *rest, is_leaf=lambda a: isinstance(a, bu.Quantity))
 
 
-
 def euler_step(target, t: jax.typing.ArrayLike, *args):
   with bst.environ.context(t=t):
     return target(*args)

dendritex/channels/calcium.py

@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst

dendritex/channels/hyperpolarization_activated.py

@@ -4,6 +4,8 @@
 This module implements hyperpolarization-activated cation channels.
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst

dendritex/channels/leaky.py

@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Sequence, Optional
 
 import brainstate as bst

dendritex/channels/potassium.py

@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional, Sequence
 
 import brainstate as bst

dendritex/channels/potassium_calcium.py

@@ -5,6 +5,8 @@
 This module implements calcium-dependent potassium channels.
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst

dendritex/channels/sodium.py

@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst

dendritex/ions/calcium.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst

dendritex/ions/potassium.py

@@ -14,10 +14,12 @@
 # ==============================================================================
 
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
-import brainunit as bu
 import brainstate as bst
+import brainunit as bu
 
 from .._base import Ion, Channel, check_hierarchies
 

dendritex/ions/sodium.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
 
 from typing import Union, Callable, Optional
 
 import brainstate as bst
 import brainunit as bu
 
-from .._base import Ion, Channel, check_hierarchies
+from .._base import Ion
 
 __all__ = [
   'Sodium',
@@ -57,4 +58,3 @@ def __init__(
     )
     self.E = bst.init.param(E, self.varshape, allow_none=False)
     self.C = bst.init.param(C, self.varshape, allow_none=False)
-

dendritex/neurons/__init__.py

@@ -17,6 +17,4 @@
 from .hh import *
 from .hh import __all__ as _hh_all
 
-
 __all__ = _hh_all
-

dendritex/neurons/hh.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
 
 from typing import Union, Optional, Callable