Repository: PlasmaPy/PlasmaPy

Issue:
Get `@particle_input` to pass mypy's strict mode
We use `@particle_input` throughout the package, so it should be a high priority to get it to work with mypy.  In particular, we should get it to work with mypy's strict mode. 

See also #2424 & #2429.


Discussion:

Resolving patch:
diff --git a/plasmapy/particles/decorators.py b/plasmapy/particles/decorators.py
--- a/plasmapy/particles/decorators.py
+++ b/plasmapy/particles/decorators.py
@@ -2,16 +2,17 @@
 
 __all__ = ["particle_input"]
 
+
 import functools
 import inspect
 import numpy as np
 import warnings
 import wrapt
 
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, MutableMapping
 from inspect import BoundArguments
 from numbers import Integral, Real
-from typing import Any, Optional, Union
+from typing import Any, Optional, TypedDict, Union
 
 from plasmapy.particles._factory import _physical_particle_factory
 from plasmapy.particles.exceptions import (
@@ -26,6 +27,19 @@
 from plasmapy.particles.particle_collections import ParticleList, ParticleListLike
 from plasmapy.utils.exceptions import PlasmaPyDeprecationWarning
 
+
+class _CallableDataDict(TypedDict, total=False):
+    allow_custom_particles: bool
+    allow_particle_lists: bool
+    annotations: dict[str, Any]
+    any_of: Optional[Union[str, Iterable[str]]]
+    callable_: Callable[..., Any]
+    exclude: Optional[Union[str, Iterable[str]]]
+    parameters_to_process: list[str]
+    require: Optional[Union[str, Iterable[str]]]
+    signature: inspect.Signature
+
+
 _basic_particle_input_annotations = (
     Particle,  # deprecated
     ParticleLike,
@@ -43,7 +57,7 @@
 )
 
 
-def _get_annotations(callable_: Callable):
+def _get_annotations(callable_: Callable[..., Any]) -> dict[str, Any]:
     """
     Access the annotations of a callable.
 
@@ -53,10 +67,10 @@ def _get_annotations(callable_: Callable):
        `inspect.get_annotations`.
     """
     # Python 3.10: Replace this with inspect.get_annotations
-    return getattr(callable_, "__annotations__", None)
+    return getattr(callable_, "__annotations__", {})
 
 
-def _make_into_set_or_none(obj) -> Optional[set]:
+def _make_into_set_or_none(obj: Any) -> Optional[Iterable[str]]:
     """
     Return `None` if ``obj`` is `None`, and otherwise convert ``obj``
     into a `set`.
@@ -71,10 +85,10 @@ def _make_into_set_or_none(obj) -> Optional[set]:
 
 def _bind_arguments(
     wrapped_signature: inspect.Signature,
-    callable_: Callable,
-    args: Optional[tuple] = None,
-    kwargs: Optional[dict[str, Any]] = None,
-    instance=None,
+    callable_: Callable[..., Any],
+    args: Iterable[Any],
+    kwargs: MutableMapping[str, Any],
+    instance: Any = None,
 ) -> inspect.BoundArguments:
     """
     Bind the arguments provided by ``args`` and ``kwargs`` to
@@ -177,16 +191,16 @@ class _ParticleInput:
 
     def __init__(
         self,
-        callable_: Callable,
+        callable_: Callable[..., Any],
         *,
-        require: Optional[Union[str, set, list, tuple]] = None,
-        any_of: Optional[Union[str, set, list, tuple]] = None,
-        exclude: Optional[Union[str, set, list, tuple]] = None,
+        require: Optional[Union[str, Iterable[str]]] = None,
+        any_of: Optional[Union[str, Iterable[str]]] = None,
+        exclude: Optional[Union[str, Iterable[str]]] = None,
         allow_custom_particles: bool = True,
         allow_particle_lists: bool = True,
     ) -> None:
-        self._data = {}
-        self.callable_ = callable_
+        self._data: _CallableDataDict = {}
+        self.callable_: Callable[..., Any] = callable_
         self.require = require
         self.any_of = any_of
         self.exclude = exclude
@@ -194,7 +208,7 @@ def __init__(
         self.allow_particle_lists = allow_particle_lists
 
     @property
-    def callable_(self) -> Callable:
+    def callable_(self) -> Callable[..., Any]:
         """
         The callable that is being decorated.
 
@@ -202,11 +216,11 @@ def callable_(self) -> Callable:
         -------
         callable
         """
-        return self._data["callable"]
+        return self._data["callable_"]
 
     @callable_.setter
-    def callable_(self, callable_: Callable) -> None:
-        self._data["callable"] = callable_
+    def callable_(self, callable_: Callable[..., Any]) -> None:
+        self._data["callable_"] = callable_
         self._data["annotations"] = _get_annotations(callable_)
         self._data["parameters_to_process"] = self.find_parameters_to_process()
         self._data["signature"] = inspect.signature(callable_)
@@ -240,10 +254,10 @@ def annotations(self) -> dict[str, Any]:
         -------
         `dict` of `str` to `object`
         """
-        return self._data.get("annotations")
+        return self._data.get("annotations")  # type: ignore[return-value]
 
     @property
-    def require(self) -> Optional[set]:
+    def require(self) -> Optional[Iterable[str]]:
         """
         Categories that the particle must belong to.
 
@@ -254,11 +268,11 @@ def require(self) -> Optional[set]:
         return self._data["require"]
 
     @require.setter
-    def require(self, require_: Optional[Union[str, set, list, tuple]]) -> None:
+    def require(self, require_: Optional[Union[str, Iterable[str]]]) -> None:
         self._data["require"] = _make_into_set_or_none(require_)
 
     @property
-    def any_of(self) -> Optional[set]:
+    def any_of(self) -> Optional[Iterable[str]]:
         """
         Categories of which the particle must belong to at least one.
 
@@ -269,11 +283,11 @@ def any_of(self) -> Optional[set]:
         return self._data["any_of"]
 
     @any_of.setter
-    def any_of(self, any_of_: Optional[Union[str, set, list, tuple]]) -> None:
+    def any_of(self, any_of_: Optional[Union[str, Iterable[str]]]) -> None:
         self._data["any_of"] = _make_into_set_or_none(any_of_)
 
     @property
-    def exclude(self) -> Optional[set]:
+    def exclude(self) -> Optional[Iterable[str]]:
         """
         Categories that the particle cannot belong to.
 
@@ -284,7 +298,7 @@ def exclude(self) -> Optional[set]:
         return self._data["exclude"]
 
     @exclude.setter
-    def exclude(self, exclude_) -> None:
+    def exclude(self, exclude_: Optional[Union[str, Iterable[str]]]) -> None:
         self._data["exclude"] = _make_into_set_or_none(exclude_)
 
     @property
@@ -331,7 +345,9 @@ def parameters_to_process(self) -> list[str]:
         """
         return self._data["parameters_to_process"]
 
-    def verify_charge_categorization(self, particle) -> None:
+    def verify_charge_categorization(
+        self, particle: Union[Particle, CustomParticle, ParticleList]
+    ) -> None:
         """
         Raise an exception if the particle does not meet charge
         categorization criteria.
@@ -351,7 +367,7 @@ def verify_charge_categorization(self, particle) -> None:
 
         if isinstance(uncharged, Iterable):
             uncharged = any(uncharged)
-            lacks_charge_info = any(lacks_charge_info)
+            lacks_charge_info = any(lacks_charge_info)  # type: ignore[arg-type]
 
         if must_be_charged and (uncharged or must_have_charge_info):
             raise ChargeError(f"{self.callable_} can only accept charged particles.")
@@ -363,7 +379,13 @@ def verify_charge_categorization(self, particle) -> None:
             )
 
     @staticmethod
-    def category_errmsg(particle, require, exclude, any_of, callable_name) -> str:
+    def category_errmsg(
+        particle: Union[Particle, CustomParticle, ParticleList],
+        require: Optional[Union[str, Iterable[str]]],
+        exclude: Optional[Union[str, Iterable[str]]],
+        any_of: Optional[Union[str, Iterable[str]]],
+        callable_name: str,
+    ) -> str:
         """
         Return an error message for when a particle does not meet
         categorization criteria.
@@ -391,15 +413,25 @@ def category_errmsg(particle, require, exclude, any_of, callable_name) -> str:
 
         return category_errmsg
 
-    def verify_particle_categorization(self, particle) -> None:
+    def verify_particle_categorization(
+        self, particle: Union[Particle, CustomParticle, ParticleList]
+    ) -> None:
         """
         Verify that the particle meets the categorization criteria.
 
+        Parameters
+        ----------
+        particle : Particle | CustomParticle
+
         Raises
         ------
         |ParticleError|
             If the particle does not meet the categorization criteria.
 
+        Notes
+        -----
+        This method does not yet work with |ParticleList| objects.
+
         See Also
         --------
         ~plasmapy.particles.particle_class.Particle.is_category
@@ -418,7 +450,9 @@ def verify_particle_categorization(self, particle) -> None:
             )
             raise ParticleError(errmsg)
 
-    def verify_particle_name_criteria(self, parameter, particle):
+    def verify_particle_name_criteria(
+        self, parameter: str, particle: Union[Particle, CustomParticle, ParticleList]
+    ) -> None:
         """
         Check that parameters with special names meet the expected
         categorization criteria.
@@ -432,7 +466,9 @@ def verify_particle_name_criteria(self, parameter, particle):
         ):
             return
 
-        name_categorization_exception = [
+        name_categorization_exception: list[
+            tuple[str, dict[str, Optional[Union[str, Iterable[str]]]], type]
+        ] = [
             ("element", {"require": "element"}, InvalidElementError),
             ("isotope", {"require": "isotope"}, InvalidIsotopeError),
             (
@@ -449,7 +485,7 @@ def verify_particle_name_criteria(self, parameter, particle):
             meets_name_criteria = particle.is_category(**categorization)
 
             if isinstance(particle, Iterable) and not isinstance(particle, str):
-                meets_name_criteria = all(meets_name_criteria)
+                meets_name_criteria = all(meets_name_criteria)  # type: ignore[arg-type]
 
             if not meets_name_criteria:
                 raise exception(
@@ -458,7 +494,9 @@ def verify_particle_name_criteria(self, parameter, particle):
                     f"valid {parameter}."
                 )
 
-    def verify_allowed_types(self, particle):
+    def verify_allowed_types(
+        self, particle: Union[Particle, CustomParticle, ParticleList]
+    ) -> None:
         """
         Verify that the particle object contains only the allowed types
         of particles.
@@ -489,8 +527,8 @@ def process_argument(
         self,
         parameter: str,
         argument: Any,
-        Z: Optional[Integral],
-        mass_numb: Optional[Integral],
+        Z: Optional[float],
+        mass_numb: Optional[int],
     ) -> Any:
         """
         Process an argument that has an appropriate annotation.
@@ -562,7 +600,9 @@ def process_argument(
 
     parameters_to_skip = ("Z", "mass_numb")
 
-    def perform_pre_validations(self, Z, mass_numb):
+    def perform_pre_validations(
+        self, Z: Optional[float], mass_numb: Optional[int]
+    ) -> None:
         """
         Perform a variety of pre-checks on the arguments.
 
@@ -594,7 +634,10 @@ def perform_pre_validations(self, Z, mass_numb):
             )
 
     def process_arguments(
-        self, args: tuple, kwargs: dict[str, Any], instance=None
+        self,
+        args: Iterable[Any],
+        kwargs: MutableMapping[str, Any],
+        instance: Any = None,
     ) -> BoundArguments:
         """
         Process the arguments passed to the callable_ callable.
@@ -639,14 +682,14 @@ def process_arguments(
 
 
 def particle_input(
-    callable_: Optional[Callable] = None,
+    callable_: Optional[Callable[..., Any]] = None,
     *,
-    require: Optional[Union[str, set, list, tuple]] = None,
-    any_of: Optional[Union[str, set, list, tuple]] = None,
-    exclude: Optional[Union[str, set, list, tuple]] = None,
+    require: Optional[Union[str, Iterable[str]]] = None,
+    any_of: Optional[Union[str, Iterable[str]]] = None,
+    exclude: Optional[Union[str, Iterable[str]]] = None,
     allow_custom_particles: bool = True,
     allow_particle_lists: bool = True,
-) -> Callable:
+) -> Callable[..., Any]:
     r"""
     Convert |particle-like| |arguments| into particle objects.
 
@@ -918,9 +961,15 @@ def instance_method(self, particle: ParticleLike, B: u.Quantity[u.T]):
 
     @wrapt.decorator
     def wrapper(
-        callable__: Callable, instance: Any, args: tuple, kwargs: dict[str, Any]
-    ):
+        callable__: Callable[..., Any],
+        instance: Any,
+        args: Iterable[Any],
+        kwargs: MutableMapping[str, Any],
+    ) -> Callable[..., Any]:
         bound_arguments = particle_validator.process_arguments(args, kwargs, instance)
-        return callable__(*bound_arguments.args, **bound_arguments.kwargs)
+        return callable__(  # type: ignore[no-any-return]
+            *bound_arguments.args,
+            **bound_arguments.kwargs,
+        )
 
-    return wrapper(callable_)
+    return wrapper(callable_, instance=None, args=(), kwargs={})


Test patch:
diff --git a/plasmapy/particles/tests/test_decorators.py b/plasmapy/particles/tests/test_decorators.py
--- a/plasmapy/particles/tests/test_decorators.py
+++ b/plasmapy/particles/tests/test_decorators.py
@@ -1,12 +1,13 @@
+"""Tests for `plasmapy.particles.decorators`."""
+
 import astropy.constants as const
 import astropy.units as u
 import inspect
 import pytest
 
-from numbers import Real
-from typing import Optional
+from collections.abc import Callable, Iterable
+from typing import Any, Optional, Union
 
-from plasmapy.particles import ParticleList
 from plasmapy.particles.decorators import particle_input
 from plasmapy.particles.exceptions import (
     ChargeError,
@@ -17,6 +18,7 @@
     ParticleError,
 )
 from plasmapy.particles.particle_class import CustomParticle, Particle, ParticleLike
+from plasmapy.particles.particle_collections import ParticleList
 from plasmapy.utils.code_repr import call_string
 from plasmapy.utils.decorators.validators import validate_quantities
 from plasmapy.utils.exceptions import PlasmaPyDeprecationWarning
@@ -24,12 +26,12 @@
 
 @particle_input
 def function_decorated_with_particle_input(
-    a,
+    a: Any,
     particle: ParticleLike,
-    b=None,
-    Z: Optional[int] = None,
+    b: Any = None,
+    Z: Optional[float] = None,
     mass_numb: Optional[int] = None,
-) -> Particle:
+) -> Union[Particle, CustomParticle, ParticleList]:
     """
     A simple function that is decorated with `particle_input` and
     returns the particle instance corresponding to the inputs.
@@ -48,12 +50,12 @@ def function_decorated_with_particle_input(
 
 @particle_input()
 def function_decorated_with_call_of_particle_input(
-    a,
+    a: Any,
     particle: ParticleLike,
-    b=None,
-    Z: Optional[int] = None,
+    b: Any = None,
+    Z: Optional[float] = None,
     mass_numb: Optional[int] = None,
-) -> Particle:
+) -> Union[Particle, CustomParticle, ParticleList]:
     """
     A simple function that is decorated with `@particle_input()` and
     returns the Particle instance corresponding to the inputs.
@@ -70,7 +72,9 @@ def function_decorated_with_call_of_particle_input(
     return particle
 
 
-particle_input_simple_table = [
+particle_input_simple_table: list[
+    tuple[tuple[Union[int, str], ...], dict[str, Any], str]
+] = [
     ((1, "p+"), {"b": 2}, "p+"),
     ((1, "Fe"), {"mass_numb": 56, "Z": 3}, "Fe-56 3+"),
     ((1,), {"particle": "e-"}, "e-"),
@@ -86,7 +90,12 @@ def function_decorated_with_call_of_particle_input(
     ],
 )
 @pytest.mark.parametrize(("args", "kwargs", "symbol"), particle_input_simple_table)
-def test_particle_input_simple(func, args, kwargs, symbol) -> None:
+def test_particle_input_simple(
+    func: Callable[..., Any],
+    args: tuple[Union[int, str], ...],
+    kwargs: dict[str, Any],
+    symbol: str,
+) -> None:
     """
     Test that simple functions decorated by particle_input correctly
     return the correct Particle object.
@@ -127,7 +136,9 @@ def test_particle_input_simple(func, args, kwargs, symbol) -> None:
 @pytest.mark.parametrize(
     ("func", "kwargs", "expected_error"), particle_input_error_table
 )
-def test_particle_input_errors(func, kwargs, expected_error) -> None:
+def test_particle_input_errors(
+    func: Callable[..., Any], kwargs: dict[str, Any], expected_error: type
+) -> None:
     """
     Test that functions decorated with `@particle_input` raise the
     expected errors.
@@ -144,16 +155,23 @@ class ClassWithDecoratedMethods:
     """
 
     @particle_input
-    def method_decorated_with_no_parentheses(self, particle: ParticleLike) -> Particle:
-        return particle
+    def method_decorated_with_no_parentheses(
+        self, particle: ParticleLike
+    ) -> Union[Particle, CustomParticle, ParticleList]:
+        # Because run-time logic is needed, mypy is unable to infer
+        # that @particle_input is doing a type conversion here. We
+        # would need a dynamic type checker to address this case.
+        return particle  # type: ignore[return-value]
 
     @particle_input()
-    def method_decorated_with_parentheses(self, particle: ParticleLike) -> Particle:
-        return particle
+    def method_decorated_with_parentheses(
+        self, particle: ParticleLike
+    ) -> Union[Particle, CustomParticle, ParticleList]:
+        return particle  # type: ignore[return-value]
 
 
 @pytest.mark.parametrize("symbol", ["muon", "He 2+"])
-def test_particle_input_classes(symbol) -> None:
+def test_particle_input_classes(symbol: str) -> None:
     """Test that @particle_input works for instance methods."""
     expected = Particle(symbol)
     instance = ClassWithDecoratedMethods()
@@ -183,7 +201,10 @@ def function_with_no_annotations() -> None:
 
 @particle_input
 def ambiguous_keywords(
-    p1: ParticleLike, p2: ParticleLike, Z=None, mass_numb=None
+    p1: ParticleLike,
+    p2: ParticleLike,
+    Z: Optional[float] = None,
+    mass_numb: Optional[int] = None,
 ) -> None:
     """
     A trivial function with two annotated arguments plus the keyword
@@ -199,7 +220,7 @@ def ambiguous_keywords(
 
 
 @pytest.mark.parametrize(("args", "kwargs"), ambiguous_arguments)
-def test_function_with_ambiguity(args, kwargs) -> None:
+def test_function_with_ambiguity(args: tuple[str, ...], kwargs: dict[str, int]) -> None:
     """
     Test that a function decorated with particle_input that has two
     annotated arguments along with `Z` and `mass_numb` raises an
@@ -214,8 +235,10 @@ def test_optional_particle() -> None:
     particle = "He"
 
     @particle_input
-    def has_default_particle(particle: ParticleLike = particle):
-        return particle
+    def has_default_particle(
+        particle: ParticleLike = particle,
+    ) -> Union[Particle, CustomParticle, ParticleList]:
+        return particle  # type: ignore[return-value]
 
     assert has_default_particle() == Particle(particle)
     assert has_default_particle("Ne") == Particle("Ne")
@@ -264,7 +287,11 @@ def has_default_particle(particle: ParticleLike = particle):
 @pytest.mark.parametrize(
     ("categorization", "particle", "exception"), categorization_particle_exception
 )
-def test_decorator_categories(categorization, particle, exception) -> None:
+def test_decorator_categories(
+    categorization: dict[str, Union[str, Iterable[str]]],
+    particle: ParticleLike,
+    exception: type,
+) -> None:
     """
     Test that the ``require``, ``any_of``, and ``exclude`` categories
     lead to a |ParticleError| being raised when an inputted particle
@@ -272,9 +299,11 @@ def test_decorator_categories(categorization, particle, exception) -> None:
     to a |ParticleError| when the inputted particle matches the criteria.
     """
 
-    @particle_input(**categorization)
-    def decorated_function(argument: ParticleLike) -> Particle:
-        return argument
+    @particle_input(**categorization)  # type: ignore[arg-type]
+    def decorated_function(
+        argument: ParticleLike,
+    ) -> Union[Particle, CustomParticle, ParticleList]:
+        return argument  # type: ignore[return-value]
 
     if exception:
         with pytest.raises(exception):
@@ -296,7 +325,7 @@ def test_optional_particle_annotation_parameter() -> None:
 
     @particle_input
     def func_optional_particle(particle: Optional[ParticleLike]) -> Optional[Particle]:
-        return particle
+        return particle  # type: ignore[return-value]
 
     assert func_optional_particle(None) is None, (
         "The particle keyword in the particle_input decorator is set "
@@ -305,20 +334,24 @@ def func_optional_particle(particle: Optional[ParticleLike]) -> Optional[Particl
     )
 
 
-def undecorated_function(particle: ParticleLike, distance: u.Quantity[u.m]):
+def undecorated_function(
+    particle: ParticleLike, distance: u.Quantity[u.m]
+) -> tuple[ParticleLike, u.Quantity[u.m]]:
     return particle, distance
 
 
 decorator_pairs = [
     (particle_input, validate_quantities),
     (particle_input(), validate_quantities),
-    (particle_input, validate_quantities()),
-    (particle_input(), validate_quantities()),
+    (particle_input, validate_quantities()),  # type: ignore[no-untyped-call]
+    (particle_input(), validate_quantities()),  # type: ignore[no-untyped-call]
 ]
 
 
 @pytest.mark.parametrize(("decorator1", "decorator2"), decorator_pairs)
-def test_stacking_decorators(decorator1, decorator2) -> None:
+def test_stacking_decorators(
+    decorator1: Callable[..., Any], decorator2: Callable[..., Any]
+) -> None:
     """
     Test that particle_input and validate_quantities can be stacked in
     either order with or without parentheses.
@@ -338,7 +371,9 @@ def test_stacking_decorators(decorator1, decorator2) -> None:
 
 
 @pytest.mark.parametrize(("decorator1", "decorator2"), decorator_pairs)
-def test_preserving_signature_with_stacked_decorators(decorator1, decorator2) -> None:
+def test_preserving_signature_with_stacked_decorators(
+    decorator1: Callable[..., Any], decorator2: Callable[..., Any]
+) -> None:
     """
     Test that |particle_input| & |validate_quantities| preserve the
     function signature after being stacked.
@@ -362,8 +397,10 @@ def test_annotated_classmethod() -> None:
     class HasAnnotatedClassMethod:
         @classmethod
         @particle_input
-        def f(cls, particle: ParticleLike) -> Particle:
-            return particle
+        def f(
+            cls, particle: ParticleLike
+        ) -> Union[Particle, CustomParticle, ParticleList]:
+            return particle  # type: ignore[return-value]
 
     has_annotated_classmethod = HasAnnotatedClassMethod()
     assert has_annotated_classmethod.f("p+") == Particle("p+")
@@ -371,13 +408,17 @@ def f(cls, particle: ParticleLike) -> Particle:
 
 @pytest.mark.parametrize("outer_decorator", [particle_input, particle_input()])
 @pytest.mark.parametrize("inner_decorator", [particle_input, particle_input()])
-def test_self_stacked_decorator(outer_decorator, inner_decorator) -> None:
+def test_self_stacked_decorator(
+    outer_decorator: Callable[..., Any], inner_decorator: Callable[..., Any]
+) -> None:
     """Test that particle_input can be stacked with itself."""
 
     @outer_decorator
     @inner_decorator
-    def f(x, particle: ParticleLike):
-        return particle
+    def f(
+        x: Any, particle: ParticleLike
+    ) -> Union[Particle, CustomParticle, ParticleList]:
+        return particle  # type: ignore[return-value]
 
     result = f(1, "p+")
     assert result == "p+"
@@ -386,7 +427,9 @@ def f(x, particle: ParticleLike):
 
 @pytest.mark.parametrize("outer_decorator", [particle_input, particle_input()])
 @pytest.mark.parametrize("inner_decorator", [particle_input, particle_input()])
-def test_class_stacked_decorator(outer_decorator, inner_decorator) -> None:
+def test_class_stacked_decorator(
+    outer_decorator: Callable[..., Any], inner_decorator: Callable[..., Any]
+) -> None:
     """
     Test that particle_input can be stacked with itself for an
     instance method.
@@ -403,7 +446,7 @@ def __init__(self, particle: ParticleLike) -> None:
     assert isinstance(result.particle, Particle)
 
 
-validate_quantities_ = validate_quantities(
+validate_quantities_ = validate_quantities(  # type: ignore[no-untyped-call]
     T_e={"equivalencies": u.temperature_energy()}
 )
 
@@ -413,7 +456,7 @@ def test_annotated_init() -> None:
 
     class HasAnnotatedInit:
         @particle_input(require="element")
-        def __init__(self, particle: ParticleLike, ionic_fractions=None) -> None:
+        def __init__(self, particle: ParticleLike, ionic_fractions: Any = None) -> None:
             self.particle = particle
 
     x = HasAnnotatedInit("H-1", ionic_fractions=32)
@@ -444,7 +487,8 @@ def __init__(self, particle: ParticleLike, ionic_fractions=None) -> None:
     ],
 )
 def test_particle_input_with_validate_quantities(
-    outer_decorator, inner_decorator
+    outer_decorator: Callable[..., Any],
+    inner_decorator: Callable[..., Any],
 ) -> None:
     """Test that particle_input can be stacked with validate_quantities."""
 
@@ -480,12 +524,14 @@ def __init__(
 @pytest.mark.parametrize(
     ("kwargs_to_particle_input", "arg"), kwargs_to_decorator_and_args
 )
-def test_particle_input_verification(kwargs_to_particle_input, arg) -> None:
+def test_particle_input_verification(
+    kwargs_to_particle_input: dict[str, Any], arg: Any
+) -> None:
     """Test the allow_custom_particles keyword argument to particle_input."""
 
     @particle_input(**kwargs_to_particle_input)
-    def f(particle: ParticleLike):
-        return particle
+    def f(particle: ParticleLike) -> Union[Particle, CustomParticle, ParticleList]:
+        return particle  # type: ignore[return-value]
 
     with pytest.raises(ParticleError):
         f(arg)
@@ -499,11 +545,11 @@ class ParameterNamesCase:
 
     def __init__(
         self,
-        category,
-        function,
-        particles_in_category,
-        particles_not_in_category,
-        exception,
+        category: str,
+        function: Callable[..., Any],
+        particles_in_category: list[ParticleLike],
+        particles_not_in_category: list[ParticleLike],
+        exception: type,
     ) -> None:
         self.category = category
         self.function = function
@@ -516,18 +562,20 @@ def __init__(
 
 
 @particle_input
-def get_element(element: ParticleLike):
-    return element
+def get_element(element: ParticleLike) -> Union[Particle, CustomParticle, ParticleList]:
+    return element  # type: ignore[return-value]
 
 
 @particle_input
-def get_isotope(isotope: ParticleLike):
-    return isotope
+def get_isotope(isotope: ParticleLike) -> Union[Particle, CustomParticle, ParticleList]:
+    return isotope  # type: ignore[return-value]
 
 
 @particle_input
-def get_ion(ion: ParticleLike, Z: Optional[Real] = None):
-    return ion
+def get_ion(
+    ion: ParticleLike, Z: Optional[float] = None
+) -> Union[Particle, CustomParticle, ParticleList]:
+    return ion  # type: ignore[return-value]
 
 
 cases = [
@@ -564,7 +612,9 @@ class TestParticleInputParameterNames:
     categories.
     """
 
-    def test_individual_particles_not_in_category(self, case) -> None:
+    def test_individual_particles_not_in_category(
+        self, case: ParameterNamesCase
+    ) -> None:
         """
         Test that the appropriate exception is raised when the function
         is provided with individual particles that are not in the
@@ -574,7 +624,7 @@ def test_individual_particles_not_in_category(self, case) -> None:
             with pytest.raises(case.exception):
                 case.function(particle)
 
-    def test_particle_list_not_in_category(self, case) -> None:
+    def test_particle_list_not_in_category(self, case: ParameterNamesCase) -> None:
         """
         Test that the appropriate exception is raised when the function
         is provided with multiple particles at once that are all not in
@@ -583,7 +633,7 @@ def test_particle_list_not_in_category(self, case) -> None:
         with pytest.raises(case.exception):
             case.function(case.particles_not_in_category)
 
-    def test_particle_list_some_in_category(self, case) -> None:
+    def test_particle_list_some_in_category(self, case: ParameterNamesCase) -> None:
         """
         Test that the appropriate exception is raised when the function
         is provided with multiple particles at once, of which some are
@@ -595,7 +645,9 @@ def test_particle_list_some_in_category(self, case) -> None:
     # If "ionic_level" gets added as a particle category, then add
     # assertions to the following test using is_category.
 
-    def test_individual_particles_all_in_category(self, case) -> None:
+    def test_individual_particles_all_in_category(
+        self, case: ParameterNamesCase
+    ) -> None:
         """
         Test that no exception is raised when the function is provided
         with individual particles that are all in the category.
@@ -603,7 +655,7 @@ def test_individual_particles_all_in_category(self, case) -> None:
         for particle in case.particles_in_category:
             case.function(particle)
 
-    def test_particle_list_all_in_category(self, case) -> None:
+    def test_particle_list_all_in_category(self, case: ParameterNamesCase) -> None:
         """
         Test that no exception is raised when the function is provided
         with multiple particles at once which are all in the category.
@@ -628,9 +680,11 @@ def test_creating_mean_particle_for_parameter_named_ion() -> None:
 
 
 @particle_input
-def return_particle(particle: ParticleLike, Z=None, mass_numb=None):
+def return_particle(
+    particle: ParticleLike, Z: Optional[float] = None, mass_numb: Optional[int] = None
+) -> Union[Particle, CustomParticle, ParticleList]:
     """A simple function that is decorated by particle_input."""
-    return particle
+    return particle  # type: ignore[return-value]
 
 
 def test_particle_input_warning_for_integer_z_mean() -> None:
@@ -667,8 +721,10 @@ def test_particle_input_with_var_positional_arguments() -> None:
     """
 
     @particle_input
-    def function_with_var_positional_arguments(*args, particle: ParticleLike = None):
-        return args, particle
+    def function_with_var_positional_arguments(
+        *args: Any, particle: ParticleLike = None
+    ) -> tuple[tuple[Any], Particle]:
+        return args, particle  # type: ignore[return-value]
 
     args = (5, 6, 7)
     particle = "p+"
@@ -688,9 +744,9 @@ def test_particle_input_with_pos_and_var_positional_arguments() -> None:
 
     @particle_input
     def function_with_pos_and_var_positional_arguments(
-        a, *args, particle: ParticleLike = None
-    ):
-        return args, particle
+        a: Any, *args: Any, particle: ParticleLike = None
+    ) -> tuple[tuple[Any, ...], Particle]:
+        return args, particle  # type: ignore[return-value]
 
     a = 1
     args = (5, 6, 7)