adding more typehints

src/pint/gridutils.py

@@ -1,10 +1,14 @@
 """Tools for building chi-squared grids."""
+
+from ast import Dict
 import concurrent.futures
 import copy
 import multiprocessing
 import subprocess
 import sys
 
+from typing import List, Tuple, Optional, Dict
+
 import numpy as np
 
 from loguru import logger as log
@@ -15,25 +19,26 @@
     tqdm = None
 
 from astropy.utils.console import ProgressBar
+from astropy import units as u
 
 from pint import fitter
 
 __all__ = ["doonefit", "grid_chisq", "grid_chisq_derived"]
 
 
-def hostinfo():
+def hostinfo() -> str:
     return subprocess.check_output("uname -a", shell=True)
 
 
-def set_log(logger_):
+def set_log(logger_: "loguru._logger.Logger") -> None:
     global log
     log = logger_
 
 
 class WrappedFitter:
     """Worker class to compute one fit with specified parameters fixed but passing other parameters to fit_toas()"""
 
-    def __init__(self, ftr, **fitargs):
+    def __init__(self, ftr: fitter.Fitter, **fitargs) -> None:
         """Worker class to computes one fit with specified parameters fixed
 
         Parameters
@@ -45,7 +50,9 @@ def __init__(self, ftr, **fitargs):
         self.ftr = ftr
         self.fitargs = fitargs
 
-    def doonefit(self, parnames, parvalues, extraparnames=[]):
+    def doonefit(
+        self, parnames: List[str], parvalues: List[str], extraparnames: List[str] = []
+    ) -> Tuple[float, List[float]]:
         """Worker process that computes one fit with specified parameters fixed
 
         Parameters
@@ -105,7 +112,12 @@ def doonefit(self, parnames, parvalues, extraparnames=[]):
         return chi2, extraparvalues
 
 
-def doonefit(ftr, parnames, parvalues, extraparnames=[]):
+def doonefit(
+    ftr: fitter.Fitter,
+    parnames: List[str],
+    parvalues: List[u.Quantity],
+    extraparnames: List[str] = [],
+) -> Tuple[float, List[float]]:
     """Worker process that computes one fit with specified parameters fixed
 
     Parameters
@@ -153,16 +165,16 @@ def doonefit(ftr, parnames, parvalues, extraparnames=[]):
 
 
 def grid_chisq(
-    ftr,
-    parnames,
-    parvalues,
-    extraparnames=[],
-    executor=None,
-    ncpu=None,
-    chunksize=1,
-    printprogress=True,
+    ftr: fitter.Fitter,
+    parnames: List[str],
+    parvalues: List[u.Quantity],
+    extraparnames: List[str] = [],
+    executor: Optional[concurrent.futures.Executor] = None,
+    ncpu: Optional[int] = None,
+    chunksize: int = 1,
+    printprogress: bool = True,
     **fitargs,
-):
+) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
     """Compute chisq over a grid of parameters
 
     Parameters
@@ -379,17 +391,17 @@ def grid_chisq(
 
 
 def grid_chisq_derived(
-    ftr,
-    parnames,
-    parfuncs,
-    gridvalues,
-    extraparnames=[],
-    executor=None,
-    ncpu=None,
-    chunksize=1,
-    printprogress=True,
+    ftr: fitter.Fitter,
+    parnames: List[str],
+    parfuncs: List[callable],
+    gridvalues: List[np.ndarray],
+    extraparnames: List[str] = [],
+    executor: Optional[concurrent.futures.Executor] = None,
+    ncpu: Optional[int] = None,
+    chunksize: int = 1,
+    printprogress: bool = True,
     **fitargs,
-):
+) -> Tuple[np.ndarray, List[np.ndarray], Dict[str, np.ndarray]]:
     """Compute chisq over a grid of derived parameters
 
     Parameters
@@ -575,16 +587,16 @@ def grid_chisq_derived(
 
 
 def tuple_chisq(
-    ftr,
-    parnames,
-    parvalues,
-    extraparnames=[],
-    executor=None,
-    ncpu=None,
-    chunksize=1,
-    printprogress=True,
+    ftr: fitter.Fitter,
+    parnames: List[str],
+    parvalues: List[np.ndarray],
+    extraparnames: List[str] = [],
+    executor: Optional[concurrent.futures.Executor] = None,
+    ncpu: Optional[int] = None,
+    chunksize: int = 1,
+    printprogress: bool = True,
     **fitargs,
-):
+) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
     """Compute chisq over a list of parameter tuples
 
     Parameters
@@ -760,17 +772,17 @@ def tuple_chisq(
 
 
 def tuple_chisq_derived(
-    ftr,
-    parnames,
-    parfuncs,
-    parvalues,
-    extraparnames=[],
-    executor=None,
-    ncpu=None,
-    chunksize=1,
-    printprogress=True,
+    ftr: fitter.Fitter,
+    parnames: List[str],
+    parfuncs: List[callable],
+    parvalues: List[np.ndarray],
+    extraparnames: List[str] = [],
+    executor: Optional[concurrent.futures.Executor] = None,
+    ncpu: Optional[int] = None,
+    chunksize: int = 1,
+    printprogress: bool = True,
     **fitargs,
-):
+) -> Tuple[np.ndarray, List, Dict[str, np.ndarray]]:
     """Compute chisq over a list of derived parameter tuples
 
     Parameters