Repository: PlasmaPy/PlasmaPy
Pull request #535 (2018-09-02 01:16:49)

Use `Optional[Particle]` as an annotation
Addresses #530.

This PR adds the ability to use `Optional[Particle]` as an annotation which allows passing `None` as a Particle when a function is decorated by `@particle_input`.

```python
>>> from plasmapy.atomic import Particle, particle_input
>>> from typing import Optional

>>> @particle_input
... def take_optional_particle(particle: Optional[Particle]):
...     return particle

>>> take_optional_particle(None)
None
```

It also works with collections!

```python
>>> @particle_input
... def take_optional_particles(particles: (Optional[Particle], Particle)):
...     return particles

>>> take_optional_particles((None, 'He'))
(None, Particle("He"))

pep8speaks: Hello @ritiek! Thanks for submitting your pull request.

 - In the file [`plasmapy/atomic/particle_input.py`](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/particle_input.py), following are the PEP8 issues :

> [Line 278:100](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/particle_input.py#L278): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (122 > 99 characters)
> [Line 350:21](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/particle_input.py#L350): [E303](https://duckduckgo.com/?q=pep8%20E303) too many blank lines (2)
> [Line 355:88](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/particle_input.py#L355): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ','

 - In the file [`plasmapy/atomic/tests/test_particle_input.py`](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/tests/test_particle_input.py), following are the PEP8 issues :

> [Line 402:100](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/tests/test_particle_input.py#L402): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (117 > 99 characters)
> [Line 433:100](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/tests/test_particle_input.py#L433): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (101 > 99 characters)
> [Line 466:100](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/tests/test_particle_input.py#L466): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (119 > 99 characters)
> [Line 492:100](https://github.com/PlasmaPy/PlasmaPy/blob/7e6b51a19095abdfd34e0b511c58b025c8baa12d/plasmapy/atomic/tests/test_particle_input.py#L492): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (123 > 99 characters)

PlasmaPy follows the [PEP8 style guide](https://www.python.org/dev/peps/pep-0008/), which may be checked locally using [pycodestyle](http://pycodestyle.pycqa.org/en/latest/intro.html)

Patch:
diff --git a/plasmapy/atomic/particle_input.py b/plasmapy/atomic/particle_input.py
index dcc9227b..5b0de132 100644
--- a/plasmapy/atomic/particle_input.py
+++ b/plasmapy/atomic/particle_input.py
@@ -6,7 +6,7 @@
 
 import functools
 import inspect
-from typing import Callable, Union, Any, Set, List, Tuple
+from typing import Callable, Union, Any, Set, List, Tuple, Optional
 import numbers
 
 from .particle_class import Particle
@@ -22,6 +22,7 @@
     "particle_input",
 ]
 
+
 def _particle_errmsg(argname: str,
                      argval: str,
                      Z: int = None,
@@ -253,6 +254,8 @@ def wrapper(*args, **kwargs):
             args_to_become_particles = []
             for argname in annotations.keys():
                 if isinstance(annotations[argname], tuple):
+                    if argname == 'return':
+                        continue
                     annotated_argnames = annotations[argname]
                     expected_params = len(annotated_argnames)
                     received_params = len(arguments[argname])
@@ -273,7 +276,9 @@ def wrapper(*args, **kwargs):
                     annotated_argnames = (annotations[argname],)
 
                 for annotated_argname in annotated_argnames:
-                    if annotated_argname is Particle and argname != 'return':
+                    is_particle = annotated_argname is Particle or \
+                                  annotated_argname is Optional[Particle]
+                    if is_particle and argname != 'return':
                         args_to_become_particles.append(argname)
 
             if not args_to_become_particles:
@@ -344,15 +349,24 @@ def wrapper(*args, **kwargs):
                     # Occasionally there will be functions where it will be
                     # useful to allow None as an argument.
 
-                    if none_shall_pass and argval is None:
-                        new_kwargs[argname] = None
-                        continue
+                    # In case annotations[argname] is a collection (which looks
+                    # like (Particle, Optional[Particle], ...) or [Particle])
+                    if isinstance(annotations[argname], tuple):
+                        optional_particle = annotations[argname][pos] is Optional[Particle]
+                    elif isinstance(annotations[argname], list):
+                        optional_particle = annotations[argname] == [Optional[Particle], ]
+                    else:
+                        # Otherwise annotations[argname] must be a Particle itself
+                        optional_particle = annotations[argname] is Optional[Particle]
 
-                    params = (argval, Z, mass_numb)
-                    particle = get_particle(argname,
-                                            params,
-                                            already_particle,
-                                            funcname)
+                    if (optional_particle or none_shall_pass) and argval is None:
+                        particle = None
+                    else:
+                        params = (argval, Z, mass_numb)
+                        particle = get_particle(argname,
+                                                params,
+                                                already_particle,
+                                                funcname)
 
                     if isinstance(raw_argval, (tuple, list)):
                         # If passed argument is a tuple or list, keep
@@ -443,7 +457,6 @@ def get_particle(argname, params, already_particle, funcname):
 
         return particle
 
-
     # The following code allows the decorator to be used either with or
     # without arguments.  This allows us to invoke the decorator either
     # as `@particle_input` or as `@particle_input()`, where the latter


Test patch:
diff --git a/plasmapy/atomic/tests/test_particle_input.py b/plasmapy/atomic/tests/test_particle_input.py
index a9d7169b..9b35c234 100644
--- a/plasmapy/atomic/tests/test_particle_input.py
+++ b/plasmapy/atomic/tests/test_particle_input.py
@@ -390,30 +390,126 @@ def decorated_function(argument: Particle) -> Particle:
 
 
 def test_none_shall_pass():
-    """Tests the `none_shall_pass` keyword argument in is_particle.  If
-    `none_shall_pass=True`, then an annotated argument should allow
-    `None` to be passed through to the decorated function.  If
-    `none_shall_pass=False`, then particle_input should raise a
-    `~plasmapy.utils.AtomicError` if an annotated argument is assigned
-    the value of `None`.
-    """
+    """Tests the `none_shall_pass` keyword argument in is_particle.
+    If `none_shall_pass=True`, then an annotated argument should allow
+    `None` to be passed through to the decorated function."""
+
     @particle_input(none_shall_pass=True)
     def func_none_shall_pass(particle: Particle) -> Optional[Particle]:
         return particle
 
-    @particle_input(none_shall_pass=False)
-    def func_none_shall_not_pass(particle: Particle) -> Particle:
-        return particle
+    @particle_input(none_shall_pass=True)
+    def func_none_shall_pass_with_tuple(particles: (Particle, Particle)) -> \
+                                       (Optional[Particle], Optional[Particle]):
+        return particles
+
+    @particle_input(none_shall_pass=True)
+    def func_none_shall_pass_with_list(particles: [Particle]) -> [Optional[Particle]]:
+        return particles
 
     assert func_none_shall_pass(None) is None, \
         ("The none_shall_pass keyword in the particle_input decorator is set "
          "to True, but is not passing through None.")
 
+    assert func_none_shall_pass_with_tuple((None, None)) == (None, None), \
+        ("The none_shall_pass keyword in the particle_input decorator is set "
+         "to True, but is not passing through None.")
+
+    assert func_none_shall_pass_with_list((None, None)) == (None, None), \
+        ("The none_shall_pass keyword in the particle_input decorator is set "
+         "to True, but is not passing through None.")
+
+
+def test_none_shall_not_pass():
+    """Tests the `none_shall_pass` keyword argument in is_particle.
+    If `none_shall_pass=False`, then particle_input should raise a
+    `TypeError` if an annotated argument is assigned the value of
+    `None`."""
+
+    @particle_input(none_shall_pass=False)
+    def func_none_shall_not_pass(particle: Particle) -> Particle:
+        return particle
+
+    @particle_input(none_shall_pass=False)
+    def func_none_shall_not_pass_with_tuple(particles: (Particle, Particle)) -> \
+                                           (Particle, Particle):
+        return particles
+
+    @particle_input(none_shall_pass=False)
+    def func_none_shall_not_pass_with_list(particles: [Particle]) -> [Particle]:
+        return particles
+
     with pytest.raises(TypeError, message=(
             "The none_shall_pass keyword in the particle_input decorator is "
             "set to False, but is not raising a TypeError.")):
         func_none_shall_not_pass(None)
 
+    with pytest.raises(TypeError, message=(
+            "The none_shall_pass keyword in the particle_input decorator is "
+            "set to False, but is not raising a TypeError.")):
+        func_none_shall_not_pass_with_tuple(('He', None))
+
+    with pytest.raises(TypeError, message=(
+            "The none_shall_pass keyword in the particle_input decorator is "
+            "set to False, but is not raising a TypeError.")):
+        func_none_shall_not_pass(('He', None))
+
+
+def test_optional_particle_annotation_argname():
+    """Tests the `Optional[Particle]` annotation argument in a function
+    decorated by `@particle_input` such that the annotated argument allows
+    `None` to be passed through to the decorated function."""
+
+    @particle_input
+    def func_optional_particle(particle: Optional[Particle]) -> Optional[Particle]:
+        return particle
+
+    @particle_input
+    def func_optional_particle_with_tuple(particles: (Particle, Optional[Particle])) -> \
+                                         (Particle, Optional[Particle]):
+        return particles
+
+    @particle_input
+    def func_optional_particle_with_list(particles: [Optional[Particle]]) -> [Optional[Particle]]:
+        return particles
+
+    assert func_optional_particle(None) is None, \
+        ("The particle keyword in the particle_input decorator is set "
+         "to accept Optional[Particle], but is not passing through None.")
+
+    assert func_optional_particle_with_tuple(('He', None)) == (Particle('He'), None), \
+        ("The particles keyword in the particle_input decorator is set "
+         "to accept Optional[Particle], but is not passing through None.")
+
+    assert func_optional_particle_with_list(('He', None)) == (Particle('He'), None), \
+        ("The particles keyword in the particle_input decorator is set "
+         "to accept Optional[Particle], but is not passing through None.")
+
+
+def test_not_optional_particle_annotation_argname():
+    """Tests the `Optional[Particle]` annotation argument in a function
+    decorated by `@particle_input` such that the annotated argument does
+    not allows `None` to be passed through to the decorated function."""
+
+    @particle_input
+    def func_not_optional_particle_with_tuple(particles: (Particle, Optional[Particle])) -> \
+                                             (Particle, Optional[Particle]):
+        return particles
+
+    @particle_input
+    def func_not_optional_particle_with_list(particles: [Particle]) -> [Particle]:
+        return particles
+
+    with pytest.raises(TypeError, message=(
+            "The particle keyword in the particle_input decorator received a "
+            "None instead of a Particle, but is not raising a TypeError.")):
+        func_not_optional_particle_with_tuple((None, 'He'))
+
+    with pytest.raises(TypeError, message=(
+            "The particle keyword in the particle_input decorator received a "
+            "None instead of a Particle, but is not raising a TypeError.")):
+        func_not_optional_particle_with_list(('He', None))
+
 
 # TODO: The following tests might be able to be cleaned up and/or
 # further parametrized since there's a fair bit of repetition.