Refactor/modernize type hints

.flake8

@@ -1,6 +1,8 @@
 [flake8]
 max-line-length = 120
 
+type-checking-pydantic-enabled = true
+
 ignore =
         # E203 is not PEP 8 compliant. https://github.com/PyCQA/pycodestyle/issues/373
         E203

.pre-commit-config.yaml

@@ -24,6 +24,8 @@ repos:
   hooks:
   -   id: flake8
       args: [--max-line-length=120]
+      additional_dependencies:
+        - flake8-modern-annotations
 
 - repo: https://github.com/asottile/blacken-docs
   rev: v1.12.1

.readthedocs.yml

@@ -3,7 +3,7 @@ version: 2
 build:
   os: "ubuntu-22.04"
   tools:
-    python: "3.9"
+    python: "3.10"
 
 sphinx:
     configuration: docs/conf.py

docs/conf.py

@@ -101,8 +101,11 @@ def __getattr__(cls, name):
     "sphinxcontrib.autodoc_pydantic",
     "sphinx_design",
     "sphinx_copybutton",
+    "sphinx_resolve_py_references",
 ]
 
+autodoc2_packages = ["../libensemble"]
+
 spelling_word_list_filename = "spelling_wordlist.txt"
 spelling_ignore_pypi_package_names = True
 spelling_ignore_acronyms = True

docs/function_guides/generator.rst

@@ -143,8 +143,8 @@ Implementing functions from the above class is relatively simple:
 
     .. tab-item:: send
 
-        .. currentmodule:: libensemble.tools.persistent_support.PersistentSupport
-        .. autofunction:: send
+        .. autoclass:: libensemble.tools.persistent_support.PersistentSupport
+            :members: send
 
         This function call typically resembles::
 
@@ -154,8 +154,8 @@ Implementing functions from the above class is relatively simple:
 
     .. tab-item:: recv
 
-        .. currentmodule:: libensemble.tools.persistent_support.PersistentSupport
-        .. autofunction:: recv
+        .. autoclass:: libensemble.tools.persistent_support.PersistentSupport
+            :members: recv
 
         This function call typically resembles::
 
@@ -170,8 +170,8 @@ Implementing functions from the above class is relatively simple:
 
     .. tab-item:: send_recv
 
-        .. currentmodule:: libensemble.tools.persistent_support.PersistentSupport
-        .. autofunction:: send_recv
+        .. autoclass:: libensemble.tools.persistent_support.PersistentSupport
+            :members: send_recv
 
         This function performs both of the previous functions in a single statement. Its
         usage typically resembles::

docs/programming_libE.rst

@@ -1,8 +1,6 @@
 Constructing Workflows
 ======================
 
-We now give greater detail in programming with libEnsemble.
-
 .. toctree::
    :maxdepth: 2
    :caption: The Basics

docs/requirements.txt

@@ -6,3 +6,4 @@ sphinx-design
 numpy
 sphinx_rtd_theme>1
 sphinx-copybutton
+sphinx_resolve_py_references

libensemble/alloc_funcs/give_sim_work_first.py

@@ -1,5 +1,4 @@
 import time
-from typing import Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -15,7 +14,7 @@ def give_sim_work_first(
     alloc_specs: dict,
     persis_info: dict,
     libE_info: dict,
-) -> Tuple[dict]:
+) -> tuple[dict]:
     """
     Decide what should be given to workers. This allocation function gives any
     available simulation work first, and only when all simulations are

libensemble/ensemble.py

@@ -1,7 +1,6 @@
 import importlib
 import json
 import logging
-from typing import Optional
 
 import numpy.typing as npt
 import tomli
@@ -12,8 +11,8 @@
 from libensemble.specs import AllocSpecs, ExitCriteria, GenSpecs, LibeSpecs, SimSpecs
 from libensemble.tools import add_unique_random_streams
 from libensemble.tools import parse_args as parse_args_f
-from libensemble.tools.parse_args import mpi_init
 from libensemble.tools import save_libE_output
+from libensemble.tools.parse_args import mpi_init
 from libensemble.utils.misc import specs_dump
 
 ATTR_ERR_MSG = 'Unable to load "{}". Is the function or submodule correctly named?'
@@ -52,27 +51,36 @@ class Ensemble:
             from libensemble.sim_funcs.simple_sim import norm_eval
             from libensemble.specs import ExitCriteria, GenSpecs, LibeSpecs, SimSpecs
 
-            libE_specs = LibeSpecs(nworkers=4)
-            sampling = Ensemble(libE_specs=libE_specs)
-            sampling.sim_specs = SimSpecs(
-                sim_f=norm_eval,
-                inputs=["x"],
-                outputs=[("f", float)],
-            )
-            sampling.gen_specs = GenSpecs(
-                gen_f=latin_hypercube_sample,
-                outputs=[("x", float, (1,))],
-                user={
-                    "gen_batch_size": 50,
-                    "lb": np.array([-3]),
-                    "ub": np.array([3]),
-                },
-            )
+            if __name__ == "__main__":
 
-            sampling.add_random_streams()
-            sampling.exit_criteria = ExitCriteria(sim_max=100)
+                libE_specs = LibeSpecs(nworkers=4)
 
-            if __name__ == "__main__":
+                sim_specs = SimSpecs(
+                    sim_f=norm_eval,
+                    inputs=["x"],
+                    outputs=[("f", float)],
+                )
+
+                gen_specs = GenSpecs(
+                    gen_f=latin_hypercube_sample,
+                    outputs=[("x", float, (1,))],
+                    user={
+                        "gen_batch_size": 50,
+                        "lb": np.array([-3]),
+                        "ub": np.array([3]),
+                    },
+                )
+
+                exit_criteria = ExitCriteria(sim_max=100)
+
+                sampling = Ensemble(
+                    libE_specs=libE_specs,
+                    sim_specs=sim_specs,
+                    gen_specs=gen_specs,
+                    exit_criteria=exit_criteria,
+                )
+
+                sampling.add_random_streams()
                 sampling.run()
                 sampling.save_output(__file__)
 
@@ -270,15 +278,15 @@ class Ensemble:
 
     def __init__(
         self,
-        sim_specs: Optional[SimSpecs] = SimSpecs(),
-        gen_specs: Optional[GenSpecs] = GenSpecs(),
-        exit_criteria: Optional[ExitCriteria] = {},
-        libE_specs: Optional[LibeSpecs] = LibeSpecs(),
-        alloc_specs: Optional[AllocSpecs] = AllocSpecs(),
-        persis_info: Optional[dict] = {},
-        executor: Optional[Executor] = None,
-        H0: Optional[npt.NDArray] = None,
-        parse_args: Optional[bool] = False,
+        sim_specs: SimSpecs | None = SimSpecs(),
+        gen_specs: GenSpecs | None = GenSpecs(),
+        exit_criteria: ExitCriteria | None = {},
+        libE_specs: LibeSpecs | None = LibeSpecs(),
+        alloc_specs: AllocSpecs | None = AllocSpecs(),
+        persis_info: dict | None = {},
+        executor: Executor | None = None,
+        H0: npt.NDArray | None = None,
+        parse_args: bool | None = False,
     ):
         self.sim_specs = sim_specs
         self.gen_specs = gen_specs
@@ -333,7 +341,7 @@ def libE_specs(self) -> LibeSpecs:
         return self._libE_specs
 
     @libE_specs.setter
-    def libE_specs(self, new_specs):
+    def libE_specs(self, new_specs: LibeSpecs | dict):
         # We need to deal with libE_specs being specified as dict or class, and
         #   "not" overwrite the internal libE_specs["comms"].
 

libensemble/executors/balsam_executor.py

@@ -78,7 +78,6 @@
 import logging
 import os
 import time
-from typing import Any, Dict, List, Optional, Union
 
 from balsam import util
 
@@ -106,9 +105,9 @@
 
     def __init__(
         self,
-        app: Optional[Application] = None,
+        app: Application | None = None,
         app_args: dict = None,
-        workdir: Optional[str] = None,
+        workdir: str | None = None,
         stdout: str = None,
         stderr: str = None,
         workerid: int = None,
@@ -122,7 +121,7 @@
         # May want to override workdir with Balsam value when it exists
         Task.__init__(self, app, app_args, workdir, stdout, stderr, workerid)
 
-    def _get_time_since_balsam_submit(self) -> Union[int, float]:
+    def _get_time_since_balsam_submit(self) -> float:
         """Return time since balsam task entered ``RUNNING`` state"""
         event_query = EventLog.objects.filter(job_id=self.process.id, to_state="RUNNING")
         if not len(event_query):
@@ -146,7 +145,7 @@
         if self.total_time is None:
             self.total_time = time.time() - self.submit_time
 
-    def _set_complete(self, dry_run: bool = False) -> None:
+    def _set_complete(self, dry_run: bool = False):
         """Set task as complete"""
         self.finished = True
         if dry_run:
@@ -169,7 +168,7 @@
 
         logger.info(f"Task {self.name} ended with state {self.state}")
 
-    def poll(self) -> None:
+    def poll(self):
         """Polls and updates the status attributes of the supplied task. Requests
         Job information from Balsam service."""
         if self.dry_run:
@@ -203,7 +202,7 @@
             self.state = "FAILED"
             self._set_complete()
 
-    def wait(self, timeout: Optional[int] = None) -> None:
+    def wait(self, timeout: float = None):
         """Waits on completion of the task or raises ``TimeoutExpired``.
 
         Status attributes of task are updated on completion.
@@ -280,10 +279,10 @@
     def register_app(
         self,
         BalsamApp: ApplicationDefinition,
-        app_name: Optional[str] = None,
-        calc_type: Optional[str] = None,
-        desc: str = None,
-        precedent: Optional[str] = None,
+        app_name: str | None = None,
+        calc_type: str | None = None,
+        desc: str | None = None,
+        precedent: str | None = None,
     ) -> None:
         """Registers a Balsam ``ApplicationDefinition`` to libEnsemble. This class
         instance *must* have a ``site`` and ``command_template`` specified. See
@@ -331,9 +330,9 @@
         job_mode: str = "mpi",
         queue: str = "local",
         project: str = "local",
-        optional_params: Dict[Any, Any] = {},
-        filter_tags: Dict[Any, Any] = {},
-        partitions: List[Any] = [],
+        optional_params: dict = {},
+        filter_tags: dict = {},
+        partitions: list[dict] = [],
     ) -> BatchJob:
         """
         Submits a Balsam ``BatchJob`` machine allocation request to Balsam.
@@ -435,14 +434,14 @@
 
     def submit(
         self,
-        calc_type: Optional[str] = None,
-        app_name: Optional[str] = None,
+        calc_type: str | None = None,
+        app_name: str | None = None,
         app_args: dict = None,
         num_procs: int = None,
         num_nodes: int = None,
         procs_per_node: int = None,
         max_tasks_per_node: int = None,
-        machinefile: Optional[str] = None,
+        machinefile: str | None = None,
         gpus_per_rank: int = 0,
         transfers: dict = {},
         workdir: str = "",