Typing Utilities (#265)

* remove ignores

* no drmaa stubs

* fix typing

* lint

* Update CHANGELOG.rst

Author: patricktnast

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**2.1.13 - 06/11/25**
+
+  - Type-hinting: Fix mypy errors in utilities.py
+
 **2.1.12 - 06/11/25**
 
   - Move deprecated rtd for redis

pyproject.toml

@@ -47,7 +47,6 @@ exclude = [
     "src/vivarium_cluster_tools/psimulate/runner.py",
     "src/vivarium_cluster_tools/vipin/cli.py",
     "src/vivarium_cluster_tools/vipin/perf_counters.py",
-    "src/vivarium_cluster_tools/utilities.py",
     "tests/psimulate/cluster/test_interface.py",
     "tests/psimulate/redis_dbs/test_registry.py",
     "tests/psimulate/results/test_processing.py",
@@ -56,12 +55,11 @@ exclude = [
     "tests/psimulate/test_appending_perf_logs.py",
     "tests/psimulate/test_branches.py",
     "tests/psimulate/test_runner.py",
-    "tests/psimulate/test_utilities.py",
 ]
 
 disable_error_code = []
 
 # handle mypy errors when 3rd party packages are not typed.
 [[tool.mypy.overrides]]
-module = []
+module = ["drmaa"]
 ignore_missing_imports = true

src/vivarium_cluster_tools/utilities.py

@@ -1,4 +1,3 @@
-# mypy: ignore-errors
 """
 ================================
 vivarium_cluster_tools Utilities
@@ -15,12 +14,15 @@
 import warnings
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any
+from typing import Any, ParamSpec, TypeVar
 
 from vivarium_cluster_tools.psimulate.environment import ENV_VARIABLES
 
 NUM_ROWS_PER_CENTRAL_LOG_FILE = 100_000
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
 
 def get_cluster_name() -> str:
     """Returns the hostname"""
@@ -67,7 +69,7 @@ def backoff_and_retry(
     backoff_seconds: int | float = 30,
     num_retries: int = 3,
     log_function: Callable[[str], None] = warnings.warn,
-) -> Callable:
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
     """Adds a retry handler to the decorated function.
 
     Parameters
@@ -85,9 +87,9 @@ def backoff_and_retry(
         A function that retries the decorated function a specified number of times.
     """
 
-    def _wrap(func):
+    def _wrap(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def _wrapped(*args, **kwargs):
+        def _wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
             retries = num_retries
             while retries:
                 try:

tests/psimulate/test_utilities.py

@@ -3,12 +3,22 @@
 import time
 from pathlib import Path
 from subprocess import PIPE, Popen
+from typing import Any, TypedDict
 
 import pytest
+from pytest import FixtureRequest
 
 from vivarium_cluster_tools.utilities import backoff_and_retry, mkdir
 
 
+class MkdirParams(TypedDict, total=False):
+    """Parameters for the mkdir function."""
+
+    parents: bool
+    exists_ok: bool
+    umask: int
+
+
 @pytest.fixture(
     params=[
         ({}, None),
@@ -20,11 +30,11 @@
         ({"exists_ok": True, "parents": True}, None),
     ]
 )
-def permissions_params(request):
-    return request.param
+def permissions_params(request: FixtureRequest) -> tuple[MkdirParams, str | None]:
+    return request.param  # type: ignore[no-any-return]
 
 
-def test_mkdir_set_permissions(permissions_params: list) -> None:
+def test_mkdir_set_permissions(permissions_params: tuple[MkdirParams, str | None]) -> None:
     # Get prior umask value
     prior_umask = os.umask(0)
     os.umask(prior_umask)
@@ -36,10 +46,10 @@ def test_mkdir_set_permissions(permissions_params: list) -> None:
     parent_path = cwd / parent_dir_name
     path = parent_path / child_dir_name
 
-    mkdir_params: dict = permissions_params[0]
+    mkdir_params: MkdirParams = permissions_params[0]
     permissions: str | None = permissions_params[1] if permissions_params[1] else "drwxrwxr-x"
 
-    def test_mkdir_permissions():
+    def test_mkdir_permissions() -> None:
         mkdir(path, **mkdir_params)
         proc = Popen(
             f"ls -l | grep '{parent_dir_name}' | grep '{permissions}'",
@@ -76,18 +86,18 @@ def test_mkdir_permissions():
         assert prior_umask == os.umask(prior_umask), "umask was changed and not reset"
 
 
-def test_backoff_and_retry():
+def test_backoff_and_retry() -> None:
     class WarningCatcher:
-        def __init__(self):
-            self.caught_warnings = []
+        def __init__(self) -> None:
+            self.caught_warnings: list[str] = []
 
-        def warn(self, message, *args, **kwargs):
+        def warn(self, message: str, *args: Any, **kwargs: Any) -> None:
             self.caught_warnings.append(message)
 
     wc = WarningCatcher()
 
     @backoff_and_retry(log_function=wc.warn)
-    def successful_function():
+    def successful_function() -> bool:
         return True
 
     assert successful_function()
@@ -96,7 +106,7 @@ def successful_function():
     wc = WarningCatcher()
 
     @backoff_and_retry(backoff_seconds=0.1, num_retries=5, log_function=wc.warn)
-    def failing_function():
+    def failing_function() -> None:
         raise ValueError
 
     start = time.time()