Repository: PlasmaPy/PlasmaPy
Pull request #528 (2018-08-03 15:13:22)

Fix `@particle_input` for optional parameters
Now `@particle_input` treats parameters with default value correctly. For example, this code now works:

```python
>>> from plasmapy.atomic import particle_input, Particle

>>> @particle_input
... def make_particle(particle: Particle = "He"):
...     return particle

>>> make_particle()
Particle("He")
```

This would otherwise give a `KeyError `when calling `make_particle()`:

```python
KeyError: 'particle'
```


Patch:
diff --git a/plasmapy/atomic/particle_input.py b/plasmapy/atomic/particle_input.py
index 17aa11b1..30968a0a 100644
--- a/plasmapy/atomic/particle_input.py
+++ b/plasmapy/atomic/particle_input.py
@@ -234,8 +234,18 @@ def wrapper(*args, **kwargs):
             annotations = wrapped_function.__annotations__
             bound_args = wrapped_signature.bind(*args, **kwargs)
 
+            default_arguments = bound_args.signature.parameters
             arguments = bound_args.arguments
-            argnames = bound_args.arguments.keys()
+            argnames = bound_args.signature.parameters.keys()
+
+            # Handle optional-only arguments in function declaration
+            for default_arg in default_arguments:
+                # The argument is not contained in `arguments` if the
+                # user does not explicitly pass an optional argument.
+                # In such cases, manually add it to `arguments` with
+                # the default value of parameter.
+                if default_arg not in arguments:
+                    arguments[default_arg] = default_arguments[default_arg].default
 
             funcname = wrapped_function.__name__
 


Test patch:
diff --git a/plasmapy/atomic/tests/test_particle_input.py b/plasmapy/atomic/tests/test_particle_input.py
index 17377f9a..a9d7169b 100644
--- a/plasmapy/atomic/tests/test_particle_input.py
+++ b/plasmapy/atomic/tests/test_particle_input.py
@@ -204,7 +204,7 @@ def function_to_test_annotations(particles: Union[Tuple, List], resulting_partic
     if not returned_particle_instances:
         raise AtomicError(
             f"A function decorated by particle_input did not return "
-            f"a tuple of Particle instances for input of "
+            f"a collection of Particle instances for input of "
             f"{repr(particles)}, and instead returned"
             f"{repr(resulting_particles)}.")
 
@@ -262,13 +262,50 @@ def test_list_annotation(particles: Union[Tuple, List]):
     except Exception as exc2:
         raise AtomicError(
             f"Unable to evaluate a function decorated by particle_input"
-            f" with an annotation of (Particle, Particle) for inputs of"
+            f" with an annotation of [Particle] for inputs of"
             f" {repr(particles)}."
         ) from exc2
 
     function_to_test_annotations(particles, resulting_particles)
 
 
+class TestOptionalArgs:
+    def particle_iter(self, particles):
+        return [Particle(particle) for particle in particles]
+
+    def test_optional_particle(self):
+        particle = "He"
+
+        @particle_input
+        def optional_particle(particle: Particle = particle):
+            return particle
+
+        assert optional_particle() == Particle(particle)
+        assert optional_particle("Ne") == Particle("Ne")
+
+    def test_optional_tuple(self):
+        tuple_of_particles = ("Mg", "Al")
+
+        @particle_input
+        def optional_tuple(particles: (Particle, Particle) = tuple_of_particles):
+            return particles
+
+        function_to_test_annotations(optional_tuple(), self.particle_iter(tuple_of_particles))
+        elements = ("C", "N")
+        function_to_test_annotations(optional_tuple(elements), self.particle_iter(elements))
+
+    def test_optional_list(self):
+        list_of_particles = ("Ca", "Ne")
+
+        @particle_input
+        def optional_list(particles: [Particle] = list_of_particles):
+            return particles
+
+        function_to_test_annotations(optional_list(), self.particle_iter(list_of_particles))
+        elements = ("Na", "H", "C")
+        function_to_test_annotations(optional_list(elements), self.particle_iter(elements))
+
+
 def test_invalid_number_of_tuple_elements():
     with pytest.raises(ValueError):
         # Passed 3 elements when function only takes 2 in tuple