Repository: PlasmaPy/PlasmaPy

Issue:
Adopt ``Quantity`` type annotations that preserve units
For many of our functions, we use annotations for units:

```Python
def func(x: u.m, t: u.s) -> u.m / u.s:
    return (x / t).to(u.m / u.s)
```

These are not *type hint* annotations, since the expected value is a `Quantity` in most situations.  Instead the annotations are used by decorators like our validators and `astropy.units.quantity_input`.  Unfortunately, this means that we can't use static type checkers like [mypy](http://mypy-lang.org/) to check PlasmaPy code. 😿  Using mypy has the potential to help us find bugs, as well as errors in existing type hint annotations (#268).

Astropy now allows for [Quantity type annotations that preserve units](https://docs.astropy.org/en/stable/units/type_hints.html#preserving-units).  Both `typing.Union` and `typing.Optional` are supported.

```Python
def func(x: Quantity[u.kpc]) -> Quantity[u.m]:
    return x << u.m
```

We should migrate towards using `Quantity` annotations that preserve units in PlasmaPy.  However, before we do this, we would need to change our validators to make them work with this style of annotation.


Discussion:

Resolving patch:
diff --git a/plasmapy/utils/decorators/checks.py b/plasmapy/utils/decorators/checks.py
--- a/plasmapy/utils/decorators/checks.py
+++ b/plasmapy/utils/decorators/checks.py
@@ -217,8 +217,8 @@ def _get_value_checks(
         Returns
         -------
         Dict[str, Dict[str, bool]]
-            A complete 'checks' dictionary for checking function input arguments
-            and return.
+            A complete 'checks' dictionary for checking function input
+            arguments and return.
         """
         # initialize validation dictionary
         out_checks = {}
@@ -289,18 +289,18 @@ def _check_value(  # noqa: C901, PLR0912
         Parameters
         ----------
         arg
-            The argument to be checked
+            The argument to be checked.
 
         arg_name: str
-            The name of the argument to be checked
+            The name of the argument to be checked.
 
         arg_checks: Dict[str, bool]
-            The requested checks for the argument
+            The requested checks for the argument.
 
         Raises
         ------
         ValueError
-            raised if a check fails
+            If a check fails.
 
         """
         if arg_name == "checks_on_return":
@@ -344,8 +344,8 @@ def _check_value(  # noqa: C901, PLR0912
 
 class CheckUnits(CheckBase):
     """
-    A decorator class to 'check' — limit/control — the units of input and return
-    arguments to a function or method.
+    A decorator class to 'check' — limit/control — the units of input
+    and return arguments to a function or method.
 
     Parameters
     ----------
@@ -908,44 +908,64 @@ def _check_unit_core(  # noqa: C901, PLR0912, PLR0915
         return arg, unit, equiv, err
 
     @staticmethod
-    def _condition_target_units(targets: list, from_annotations: bool = False) -> list:
+    def _condition_target_units(
+        targets: list[Union[str, u.Unit, u.Quantity]],
+        from_annotations: bool = False,
+    ) -> list:
         """
-        From a list of target units (either as a string or astropy
-        :class:`~astropy.units.Unit` objects), return a list of conditioned
-        :class:`~astropy.units.Unit` objects.
+        From a `list` of target objects that have or represent units,
+        return a `list` of conditioned :class:`~astropy.units.Unit`
+        objects.
 
         Parameters
         ----------
-        targets: list of target units
-            list of units (either as a string or :class:`~astropy.units.Unit`)
-            to be conditioned into astropy :class:`~astropy.units.Unit` objects
-
-        from_annotations: bool
-            (Default `False`) Indicates if `targets` originated from function/method
+        targets: `list` of  `str`, `~astropy.units.Unit`, or `~astropy.units.Quantity`
+            A list containing strings representing units (e.g., ``"kg"``,
+            `~astropy.units.Unit` objects (e.g., ``u.kg``), or
+            |Quantity| objects indexed with a `~astropy.units.U nit`
+            object (e.g., ``u.Quantity[u.kg]``).
+
+        from_annotations: bool, default: `False`
+            Indicates if ``targets`` originated from function/method
             annotations versus decorator input arguments.
 
         Returns
         -------
         list:
-            list of `targets` converted into astropy
-            :class:`~astropy.units.Unit` objects
+            `list` of ``targets`` converted into
+            :class:`~astropy.units.Unit` objects.
 
         Raises
         ------
         TypeError
-            If `target` is not a valid type for :class:`~astropy.units.Unit` when
-            `from_annotations == True`,
+            If `target` is not a valid type for
+            :class:`~astropy.units.Unit` when ``from_annotations == True``,
 
         ValueError
-            If a `target` is a valid unit type but not a valid value for
-            :class:`~astropy.units.Unit`.
+            If a ``target`` is a valid unit type but not a valid value
+            for :class:`~astropy.units.Unit`.
         """
         # Note: this method does not allow for astropy physical types. This is
         #       done because we expect all use cases of CheckUnits to define the
         #       exact units desired.
-        #
+
         allowed_units = []
         for target in targets:
+            # The following two blocks are to create extract the unit from
+            # annotations of the form u.Quantity[u.m], which is an annotated
+            # alias.  The unit is stored as the first item in the __metadata__
+            # attribute and the original class is stored in the __origin__
+            # attribute.
+
+            annotation_metadata = getattr(target, "__metadata__", None)
+            annotation_original_class = getattr(target, "__origin__", None)
+
+            if (
+                annotation_original_class is u.Quantity
+                and annotation_metadata is not None
+            ):
+                target = annotation_metadata[0]  # noqa: PLW2901
+
             try:
                 target_unit = u.Unit(target)
                 allowed_units.append(target_unit)
diff --git a/plasmapy/utils/decorators/validators.py b/plasmapy/utils/decorators/validators.py
--- a/plasmapy/utils/decorators/validators.py
+++ b/plasmapy/utils/decorators/validators.py
@@ -472,7 +472,6 @@ class Foo:
             def bar(self, mass, vel):
                 return mass * vel
 
-
     Define units with function annotations::
 
         import astropy.units as u
@@ -493,6 +492,15 @@ class Foo:
             def bar(self, mass: u.g, vel: u.cm / u.s) -> u.g * u.cm / u.s:
                 return mass * vel
 
+    Define units using type hint annotations::
+
+        import astropy.units as u
+        from plasmapy.decorators import validate_quantities
+
+        @validate_quantities
+        def foo(x: u.Quantity[u.m], time: u.Quantity[u.s]) -> u.Quantity[u.m / u.s]:
+            return x / time
+
     Allow `None` values to pass::
 
         import astropy.units as u


Test patch:
diff --git a/plasmapy/utils/decorators/tests/test_validators.py b/plasmapy/utils/decorators/tests/test_validators.py
--- a/plasmapy/utils/decorators/tests/test_validators.py
+++ b/plasmapy/utils/decorators/tests/test_validators.py
@@ -6,7 +6,6 @@
 import inspect
 import pytest
 
-from functools import cached_property
 from typing import Optional
 from unittest import mock
 
@@ -18,9 +17,6 @@
 )
 
 
-# ----------------------------------------------------------------------------------------
-# Test Decorator class `ValidateQuantities` and decorator `validate_quantities`
-# ----------------------------------------------------------------------------------------
 class TestValidateQuantities:
     """
     Test for decorator
@@ -62,7 +58,7 @@ def test_vq_method__get_validations(self):
         # 'output' = expected return from `_get_validations`
         # 'raises' = if `_get_validations` raises an Exception
         # 'warns' = if `_get_validations` issues a warning
-        #
+
         _cases = [
             {
                 "descr": "typical call...using 'can_be_negative'",
@@ -229,7 +225,7 @@ def test_vq_method__validate_quantity(self):
         # 'output' = expected return from `_get_validations`
         # 'raises' = if `_get_validations` raises an Exception
         # 'warns' = if `_get_validations` issues a warning
-        #
+
         _cases = [
             # typical call
             {
@@ -298,7 +294,7 @@ def test_vq_method__validate_quantity(self):
             },
         ]
 
-        # setup wrapped function
+        # set up wrapped function
         vq = ValidateQuantities()
         vq.f = self.foo
 
@@ -346,8 +342,6 @@ def test_vq_method__validate_quantity(self):
 
     def test_vq_preserves_signature(self):
         """Test `ValidateQuantities` preserves signature of wrapped function."""
-        # I'd like to directly test the @preserve_signature is used (??)
-
         wfoo = ValidateQuantities()(self.foo_anno)
         assert hasattr(wfoo, "__signature__")
         assert wfoo.__signature__ == inspect.signature(self.foo_anno)
@@ -361,7 +355,7 @@ def test_vq_called_as_decorator(self):
         # 'output' = expected return from wrapped function
         # 'raises' = if an Exception is expected to be raised
         # 'warns' = if a warning is expected to be issued
-        #
+
         _cases = [
             {
                 "descr": "clean execution",
@@ -496,7 +490,7 @@ def test_decorator_func_def(self, mock_vq_class):
         # 'output' = expected return from wrapped function
         # 'raises' = a raised Exception is expected
         # 'warns' = an issued warning is expected
-        #
+
         _cases = [
             # only argument checks
             {
@@ -535,12 +529,10 @@ def test_decorator_func_def(self, mock_vq_class):
                     #  @validate_quantities(x=check)
                     #      def foo(x):
                     #          return x
-                    #
                     wfoo = validate_quantities(**case["setup"]["validations"])(mock_foo)
                 else:
                     continue
 
-                # test
                 args = case["setup"]["args"]
                 kwargs = case["setup"]["kwargs"]
                 assert wfoo(*args, **kwargs) == case["output"]
@@ -556,7 +548,6 @@ def test_decorator_func_def(self, mock_vq_class):
                 for arg_name, validations in case["setup"]["validations"].items():
                     assert mock_vq_class.call_args[1][arg_name] == validations
 
-                # reset
                 mock_vq_class.reset_mock()
                 mock_foo.reset_mock()
 
@@ -573,29 +564,29 @@ def __init__(
             self.y = y
             self.z = z
 
-        @cached_property
+        @property
         @validate_class_attributes(expected_attributes=["x"])
         def require_x(self):
             return 0
 
-        @cached_property
+        @property
         @validate_class_attributes(expected_attributes=["x", "y"])
         def require_x_and_y(self):
             return 0
 
-        @cached_property
+        @property
         @validate_class_attributes(both_or_either_attributes=[("x", "y")])
         def require_x_or_y(self):
             return 0
 
-        @cached_property
+        @property
         @validate_class_attributes(
             expected_attributes=["x"], both_or_either_attributes=[("y", "z")]
         )
         def require_x_and_either_y_or_z(self):
             return 0
 
-        @cached_property
+        @property
         @validate_class_attributes(mutually_exclusive_attributes=[("x", "y")])
         def require_only_either_x_or_y(self):
             return 0
@@ -637,3 +628,39 @@ def test_method_errors(self, test_case_constructor_keyword_arguments):
             else:
                 with pytest.raises(ValueError):
                     getattr(test_case, method)
+
+
+def test_validate_quantities_argument_type_annotation():
+    """
+    Test that |validate_quantities| works with type hint annotations of
+    the form ``u.Quantity[u.m]`` on a function argument.
+    """
+
+    @validate_quantities
+    def f(x: u.Quantity[u.m]):
+        return x
+
+    argument = 100 * u.cm
+    expected = 1 * u.m
+    actual = f(argument)
+
+    assert u.isclose(actual, expected)
+    assert actual.unit == expected.unit
+
+
+def test_validate_quantities_return_type_annotation():
+    """
+    Test that |validate_quantities| works with type hint annotations of
+    the form ``u.Quantity[u.m]`` as a return argument.
+    """
+
+    @validate_quantities
+    def f(x) -> u.Quantity[u.m]:
+        return x
+
+    argument = 100 * u.cm
+    expected = 1 * u.m
+    actual = f(argument)
+
+    assert u.isclose(actual, expected)
+    assert actual.unit == expected.unit