Add type hint for Callable

deepmodeling · Feb 29, 2024 · 64d6079 · 64d6079
1 parent ab35653
commit 64d6079
Show file tree

Hide file tree

Showing 12 changed files with 33 additions and 13 deletions.
diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py
@@ -92,7 +92,9 @@ def share_params(self, base_class, shared_level, resume=False):
             pass
 
         def compute_input_stats(
-            self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+            self,
+            merged: Union[Callable[[], List[dict]], List[dict]],
+            path: Optional[DPPath] = None,
         ):
             """Update mean and stddev for descriptor elements."""
             raise NotImplementedError

diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py
@@ -89,7 +89,9 @@ def get_dim_emb(self) -> int:
         pass
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for DescriptorBlock elements."""
         raise NotImplementedError

diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py
@@ -149,7 +149,9 @@ def dim_emb(self):
         return self.get_dim_emb()
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         return self.se_atten.compute_input_stats(merged, path)
 

diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py
@@ -333,7 +333,9 @@ def dim_emb(self):
         return self.get_dim_emb()
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         for ii, descrpt in enumerate([self.repinit, self.repformers]):
             descrpt.compute_input_stats(merged, path)

diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py
@@ -158,7 +158,9 @@ def share_params(self, base_class, shared_level, resume=False):
             raise NotImplementedError
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         for ii, descrpt in enumerate(self.descriptor_list):

diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
@@ -281,7 +281,9 @@ def forward(
         return g1, g2, h2, rot_mat.view(-1, nloc, self.dim_emb, 3), sw
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)

diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
@@ -145,7 +145,9 @@ def dim_out(self):
         return self.sea.dim_out
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         return self.sea.compute_input_stats(merged, path)
@@ -405,7 +407,9 @@ def __getitem__(self, key):
             raise KeyError(key)
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)

diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
@@ -203,7 +203,9 @@ def dim_emb(self):
         return self.get_dim_emb()
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
@@ -179,7 +179,9 @@ def share_params(self, base_class, shared_level, resume=False):
             raise NotImplementedError
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)

diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py
@@ -138,7 +138,7 @@ def output_def(self) -> FittingOutputDef:
 
     def compute_output_stats(
         self,
-        merged: Union[Callable, List[dict]],
+        merged: Union[Callable[[], List[dict]], List[dict]],
         stat_file_path: Optional[DPPath] = None,
     ):
         raise NotImplementedError

diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py
@@ -142,7 +142,7 @@ def serialize(self) -> dict:
 
     def compute_output_stats(
         self,
-        merged: Union[Callable, List[dict]],
+        merged: Union[Callable[[], List[dict]], List[dict]],
         stat_file_path: Optional[DPPath] = None,
     ):
         if stat_file_path is not None:

diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py
@@ -165,7 +165,7 @@ def output_def(self) -> FittingOutputDef:
 
     def compute_output_stats(
         self,
-        merged: Union[Callable, List[dict]],
+        merged: Union[Callable[[], List[dict]], List[dict]],
         stat_file_path: Optional[DPPath] = None,
     ):
         raise NotImplementedError