Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve type hint annotations to plasmapy.particles #2458

Merged
merged 18 commits into from
Mar 5, 2024
Merged
Show file tree
Hide file tree
Changes from 17 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
1 change: 1 addition & 0 deletions changelog/2458.trivial.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Updated type hint annotations in `plasmapy.particles`.
15 changes: 6 additions & 9 deletions mypy.ini
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
12 changes: 7 additions & 5 deletions plasmapy/particles/_elements.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
33 changes: 23 additions & 10 deletions plasmapy/particles/_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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)
ejohnson-96 marked this conversation as resolved.
Show resolved Hide resolved


def _physical_particle_factory(
*args, **kwargs
*args: ParticleLike | u.Quantity | CustomParticle | Sequence[ParticleLike],
Copy link
Member Author

@namurphy namurphy Mar 4, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought about creating a type alias variable for all of this since it gets repeated, but if we make _physical_particle_factory into a public function, the type alias variable would need to also be public so that the documentation build doesn't break...I think.

**kwargs: ParticleLike | u.Quantity | None,
) -> Particle | CustomParticle | ParticleList:
"""
Return a representation of one or more physical particles.
Expand Down Expand Up @@ -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.")
Expand Down