feat: Add type hints to constant_concentration.py

Author: JC230903

src/pybamm/models/submodels/electrolyte_diffusion/constant_concentration.py

@@ -1,14 +1,11 @@
-#
-# Class for leading-order electrolyte diffusion employing stefan-maxwell
-#
 import pybamm
-from typing import Dict, Optional, Any, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from pybamm import ParameterValues
+from typing import Optional, Any, TYPE_CHECKING
 
 from .base_electrolyte_diffusion import BaseElectrolyteDiffusion
 
+if TYPE_CHECKING:
+    from pybamm import ParameterValues, Symbol
+
 
 class ConstantConcentration(BaseElectrolyteDiffusion):
     """Class for constant concentration of electrolyte
@@ -21,19 +18,17 @@ class ConstantConcentration(BaseElectrolyteDiffusion):
         A dictionary of options to be passed to the model.
     """
 
-    def __init__(
-        self, param: "ParameterValues", options: Optional[Dict[str, Any]] = None
-    ) -> None:
+    def __init__(self, param: "ParameterValues", options: Optional[dict[str, Any]] = None) -> None:
         super().__init__(param, options)
 
-    def get_fundamental_variables(self) -> Dict[str, pybamm.Symbol]:
+    def get_fundamental_variables(self) -> dict[str, pybamm.Symbol]:
         c_e_init: float = self.param.c_e_init
-        eps_c_e_dict: Dict[str, pybamm.Symbol] = {
+        eps_c_e_dict: dict[str, pybamm.Symbol] = {
             domain: self.param.domain_params[domain.split()[0]].epsilon_init * c_e_init
             for domain in self.options.whole_cell_domains
         }
-        variables: Dict[str, pybamm.Symbol] = (
-            self._get_standard_porosity_times_concentration_variables(eps_c_e_dict)
+        variables: dict[str, pybamm.Symbol] = self._get_standard_porosity_times_concentration_variables(
+            eps_c_e_dict
         )
         N_e: pybamm.Symbol = pybamm.FullBroadcastToEdges(
             0,
@@ -45,16 +40,12 @@ def get_fundamental_variables(self) -> Dict[str, pybamm.Symbol]:
 
         return variables
 
-    def get_coupled_variables(
-        self, variables: Dict[str, pybamm.Symbol]
-    ) -> Dict[str, pybamm.Symbol]:
-        c_e_dict: Dict[str, pybamm.Symbol] = {}
+    def get_coupled_variables(self, variables: dict[str, pybamm.Symbol]) -> dict[str, pybamm.Symbol]:
+        c_e_dict: dict[str, pybamm.Symbol] = {}
         for domain in self.options.whole_cell_domains:
             Domain = domain.capitalize()
             eps_k: pybamm.Symbol = variables[f"{Domain} porosity"]
-            eps_c_e_k: pybamm.Symbol = variables[
-                f"{Domain} porosity times concentration [mol.m-3]"
-            ]
+            eps_c_e_k: pybamm.Symbol = variables[f"{Domain} porosity times concentration [mol.m-3]"]
             c_e_k: pybamm.Symbol = eps_c_e_k / eps_k
             c_e_dict[domain] = c_e_k
 
@@ -70,7 +61,7 @@ def get_coupled_variables(
 
         return variables
 
-    def set_boundary_conditions(self, variables: Dict[str, pybamm.Symbol]) -> None:
+    def set_boundary_conditions(self, variables: dict[str, pybamm.Symbol]) -> None:
         """
         We provide boundary conditions even though the concentration is constant
         so that the gradient of the concentration has the correct shape after
@@ -86,6 +77,6 @@ def set_boundary_conditions(self, variables: Dict[str, pybamm.Symbol]) -> None:
             }
         }
 
-    def add_events_from(self, variables: Dict[str, pybamm.Symbol]) -> None:
+    def add_events_from(self, variables: dict[str, pybamm.Symbol]) -> None:
         # No event since the concentration is constant
-        pass
+        pass