typehints for hadron read functions

Author: jkuhl-uni

pyerrors/input/hadrons.py

@@ -142,7 +142,7 @@ def read_hd5(filestem: str, ens_id: str, group: str, attrs: Optional[Union[dict,
     return corr
 
 
-def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=None):
+def read_meson_hd5(path: str, filestem: str, ens_id: str, meson: str='meson_0', idl: Optional[range]=None, gammas: Optional[tuple[str, ...]]=None):
     r'''Read hadrons meson hdf5 file and extract the meson labeled 'meson'
 
     Parameters
@@ -156,7 +156,7 @@ def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=Non
     meson : str
         label of the meson to be extracted, standard value meson_0 which
         corresponds to the pseudoscalar pseudoscalar two-point function.
-    gammas : tuple of strings
+    gammas : tuple[str]
         Instrad of a meson label one can also provide a tuple of two strings
         indicating the gamma matrices at sink and source (gamma_snk, gamma_src).
         ("Gamma5", "Gamma5") corresponds to the pseudoscalar pseudoscalar
@@ -181,7 +181,7 @@ def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=Non
                     part="real")
 
 
-def _extract_real_arrays(path, files, tree, keys):
+def _extract_real_arrays(path: str, files: list[str], tree: str, keys: list[str]):
     corr_data = {}
     for key in keys:
         corr_data[key] = []
@@ -199,7 +199,7 @@ def _extract_real_arrays(path, files, tree, keys):
     return corr_data
 
 
-def extract_t0_hd5(path, filestem, ens_id, obs='Clover energy density', fit_range=5, idl=None, **kwargs):
+def extract_t0_hd5(path: str, filestem: str, ens_id: str, obs='Clover energy density', fit_range: int=5, idl: Optional[range]=None, **kwargs):
     r'''Read hadrons FlowObservables hdf5 file and extract t0
 
     Parameters
@@ -247,7 +247,7 @@ def extract_t0_hd5(path, filestem, ens_id, obs='Clover energy density', fit_rang
     return fit_t0(t2E_dict, fit_range, plot_fit=kwargs.get('plot_fit'))
 
 
-def read_DistillationContraction_hd5(path, ens_id, diagrams=["direct"], idl=None):
+def read_DistillationContraction_hd5(path: str, ens_id: str, diagrams: list[str]=["direct"], idl: Optional[range]=None):
     """Read hadrons DistillationContraction hdf5 files in given directory structure
 
     Parameters
@@ -256,7 +256,7 @@ def read_DistillationContraction_hd5(path, ens_id, diagrams=["direct"], idl=None
         path to the directories to read
     ens_id : str
         name of the ensemble, required for internal bookkeeping
-    diagrams : list
+    diagrams : list[str]
         List of strings of the diagrams to extract, e.g. ["direct", "box", "cross"].
     idl : range
         If specified only configurations in the given range are read in.
@@ -384,7 +384,7 @@ def __array_finalize__(self, obj):
         self.mom_out = getattr(obj, 'mom_out', None)
 
 
-def read_ExternalLeg_hd5(path, filestem, ens_id, idl=None):
+def read_ExternalLeg_hd5(path: str, filestem: str, ens_id: str, idl: Optional[range]=None):
     """Read hadrons ExternalLeg hdf5 file and output an array of CObs
 
     Parameters
@@ -429,7 +429,7 @@ def read_ExternalLeg_hd5(path, filestem, ens_id, idl=None):
     return Npr_matrix(matrix, mom_in=mom)
 
 
-def read_Bilinear_hd5(path, filestem, ens_id, idl=None):
+def read_Bilinear_hd5(path: str, filestem: str, ens_id: str, idl: Optional[range]=None):
     """Read hadrons Bilinear hdf5 file and output an array of CObs
 
     Parameters
@@ -488,7 +488,7 @@ def read_Bilinear_hd5(path, filestem, ens_id, idl=None):
     return result_dict
 
 
-def read_Fourquark_hd5(path, filestem, ens_id, idl=None, vertices=["VA", "AV"]):
+def read_Fourquark_hd5(path: str, filestem: str, ens_id: str, idl: Optional[range]=None, vertices: list[str]=["VA", "AV"]):
     """Read hadrons FourquarkFullyConnected hdf5 file and output an array of CObs
 
     Parameters
@@ -574,7 +574,7 @@ def read_Fourquark_hd5(path, filestem, ens_id, idl=None, vertices=["VA", "AV"]):
     return result_dict
 
 
-def _get_lorentz_names(name):
+def _get_lorentz_names(name: str):
     lorentz_index = ['X', 'Y', 'Z', 'T']
 
     res = []