Improve type hint annotations to plasmapy.particles (#2458)

* Use and adjust monkeytype on atomic.py

* Add monkeytype file to .gitignore

* Remove unintentionally added monkeytype file

* Update type hints for atomic.py

* Type hints for particle_class.py

* Stop ignoring mypy errors in atomic.py, particle_class.py, & tests

* Update types in particle_class

* Hide some mypy errors again

* Update types in test_particle_class

* Add changelog entry

* Type hinting in _factory.py

* Type hints for _elements.py & atomic.py

* Update import

* Remove duplicate entry

* Revert a change in a comment

* Remove io.StringIO as type hint

...since it was breaking the documentation build.

changelog/2458.trivial.rst

@@ -0,0 +1 @@
+Updated type hint annotations in `plasmapy.particles`.

mypy.ini

@@ -51,6 +51,9 @@ ignore_missing_imports = true
 [mypy-pkg_resources.*]
 ignore_missing_imports = true
 
+[mypy-requests.*]
+ignore_missing_imports = true
+
 [mypy-scipy.*]
 ignore_missing_imports = true
 
@@ -312,10 +315,7 @@ disable_error_code = attr-defined,no-untyped-def
 disable_error_code = attr-defined,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.particles._elements]
-disable_error_code = arg-type,no-untyped-def
-
-[mypy-plasmapy.particles._factory]
-disable_error_code = no-any-return,no-untyped-def,operator,type-arg
+disable_error_code = arg-type,index,no-untyped-def
 
 [mypy-plasmapy.particles._isotopes]
 disable_error_code = arg-type,no-untyped-def
@@ -326,9 +326,6 @@ disable_error_code = arg-type,assignment,no-any-return,no-untyped-def,return-val
 [mypy-plasmapy.particles._special_particles]
 disable_error_code = type-arg,var-annotated
 
-[mypy-plasmapy.particles.atomic]
-disable_error_code = misc,no-any-return,no-untyped-call,no-untyped-def,union-attr
-
 [mypy-plasmapy.particles.ionization_state]
 disable_error_code = arg-type,misc,no-any-return,no-untyped-call,no-untyped-def,operator,return-value,type-arg
 
@@ -339,7 +336,7 @@ disable_error_code = assignment,has-type,index,misc,no-any-return,no-untyped-cal
 disable_error_code = no-untyped-def,union-attr
 
 [mypy-plasmapy.particles.particle_class]
-disable_error_code = arg-type,assignment,attr-defined,has-type,index,misc,no-any-return,no-redef,no-untyped-def,return-value,type-arg,valid-type,var-annotated
+disable_error_code = arg-type,assignment,attr-defined,has-type,index,misc,no-any-return,no-redef,no-untyped-call,no-untyped-def,return-value,type-arg,valid-type,var-annotated
 
 [mypy-plasmapy.particles.particle_collections]
 disable_error_code = no-any-return,no-untyped-call,no-untyped-def,override,type-arg,valid-type,var-annotated
@@ -465,7 +462,7 @@ disable_error_code = assignment,attr-defined,no-untyped-call,no-untyped-def,var-
 disable_error_code = arg-type,attr-defined,no-any-return,no-untyped-def,type-arg,union-attr
 
 [mypy-plasmapy.utils.data.downloader]
-disable_error_code = import-untyped,no-untyped-def
+disable_error_code = no-untyped-def
 
 [mypy-plasmapy.utils.data.tests.test_downloader]
 disable_error_code = no-untyped-call,no-untyped-def

plasmapy/particles/_elements.py

@@ -48,16 +48,18 @@ def element_obj_hook(obj):
 #    json.dump(_Elements, f, default=plasma_default, indent=2)
 
 
-data_about_elements = json.loads(
+data_about_elements: dict[str, str | int | u.Quantity[u.u]] = json.loads(
     pkgutil.get_data("plasmapy", "particles/data/elements.json"),
     object_hook=element_obj_hook,
 )
 
 
-atomic_numbers_to_symbols = {
-    elemdict["atomic number"]: symb for (symb, elemdict) in data_about_elements.items()
+atomic_numbers_to_symbols: dict[int, str] = {
+    elemdict["atomic number"]: symb  # type: ignore[misc]
+    for (symb, elemdict) in data_about_elements.items()
 }
 
-element_names_to_symbols = {
-    elemdict["element name"]: symb for (symb, elemdict) in data_about_elements.items()
+element_names_to_symbols: dict[str, int] = {
+    elemdict["element name"]: symb  # type: ignore[misc]
+    for (symb, elemdict) in data_about_elements.items()
 }

plasmapy/particles/_factory.py

@@ -7,19 +7,24 @@
 __all__: list[str] = []
 
 import contextlib
+from collections.abc import Callable, Sequence
 from numbers import Integral, Real
-from typing import Any
 
 import astropy.units as u
 from astropy.constants import m_e
 
 from plasmapy.particles.exceptions import ChargeError, InvalidParticleError
-from plasmapy.particles.particle_class import CustomParticle, Particle, ParticleLike
+from plasmapy.particles.particle_class import (
+    CustomParticle,
+    Particle,
+    ParticleLike,
+)
 from plasmapy.particles.particle_collections import ParticleList
 
 
 def _generate_particle_factory_error_message(
-    args: tuple, kwargs: dict[str, Any]
+    args: ParticleLike | u.Quantity | CustomParticle | Sequence[ParticleLike],
+    kwargs: dict[str, object],
 ) -> str:
     """Compose an error message for invalid particles."""
 
@@ -47,7 +52,7 @@ def _make_custom_particle_with_real_charge_number(
     mass_numb: int | None = None,
     symbol: str | None = None,
     Z: float | None = None,
-):
+) -> CustomParticle:
     """
     Create a |CustomParticle| for mean or composite ions.
 
@@ -92,26 +97,34 @@ def _make_custom_particle_with_real_charge_number(
         # Add tests if this function becomes part of public API
         raise InvalidParticleError("Cannot create CustomParticle.")
 
-    if Z > base_particle.atomic_number:
+    if Z is not None and Z > base_particle.atomic_number:
         raise ChargeError("The charge number cannot exceed the atomic number.")
 
     mass = base_particle.mass - m_e * Z
     return CustomParticle(mass=mass, Z=Z, symbol=symbol)
 
 
-_particle_constructors = (
+_particle_constructors: tuple[
+    type
+    | Callable[
+        [ParticleLike | u.Quantity | CustomParticle | Sequence[ParticleLike]],
+        object,
+    ],
+    ...,
+] = (
     Particle,
     CustomParticle,
     CustomParticle._from_quantities,  # noqa: SLF001
     ParticleList,
     _make_custom_particle_with_real_charge_number,
 )
 
-_particle_types = (Particle, CustomParticle, ParticleList)
+_particle_types: tuple[type, ...] = (Particle, CustomParticle, ParticleList)
 
 
 def _physical_particle_factory(
-    *args, **kwargs
+    *args: ParticleLike | u.Quantity | CustomParticle | Sequence[ParticleLike],
+    **kwargs: ParticleLike | u.Quantity | None,
 ) -> Particle | CustomParticle | ParticleList:
     """
     Return a representation of one or more physical particles.
@@ -182,14 +195,14 @@ def _physical_particle_factory(
             kwargs.pop(parameter)
 
     if len(args) == 1 and not kwargs and isinstance(args[0], _particle_types):
-        return args[0]
+        return args[0]  # type: ignore[return-value]
 
     if not args and not kwargs:
         raise TypeError("Particle information has not been provided.")
 
     for constructor in _particle_constructors:
         with contextlib.suppress(ChargeError, InvalidParticleError, TypeError):
-            return constructor(*args, **kwargs)
+            return constructor(*args, **kwargs)  # type: ignore[return-value]
 
     if args and not isinstance(args[0], str | Integral | u.Quantity):
         raise TypeError("Invalid type for particle.")