Repository: PlasmaPy/PlasmaPy
Pull request #1025 (2021-02-20 04:56:58)

Use yield in __iter__ for IonizationState and IonizationStateCollection
Fixes the failing test contributed by @StanczakDominik in #1021.  This makes `IonizationState` and `IonizationStateCollection` multiply iterable by using `yield`.

codecov[bot]: # [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=h1) Report
> Merging [#1025](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=desc) (3ba5a5c) into [master](https://codecov.io/gh/PlasmaPy/PlasmaPy/commit/4196ab782056646630aadb8fc31981576d488835?el=desc) (4196ab7) will **decrease** coverage by `0.79%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025/graphs/tree.svg?width=650&height=150&src=pr&token=LzZ5mnXVad)](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1025      +/-   ##
==========================================
- Coverage   96.67%   95.87%   -0.80%     
==========================================
  Files          64       64              
  Lines        6279     6262      -17     
==========================================
- Hits         6070     6004      -66     
- Misses        209      258      +49     
```


| [Impacted Files](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [plasmapy/particles/ionization\_state.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025/diff?src=pr&el=tree#diff-cGxhc21hcHkvcGFydGljbGVzL2lvbml6YXRpb25fc3RhdGUucHk=) | `93.52% <100.00%> (-0.17%)` | :arrow_down: |
| [plasmapy/particles/ionization\_state\_collection.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025/diff?src=pr&el=tree#diff-cGxhc21hcHkvcGFydGljbGVzL2lvbml6YXRpb25fc3RhdGVfY29sbGVjdGlvbi5weQ==) | `91.10% <100.00%> (-0.22%)` | :arrow_down: |
| [plasmapy/simulation/particle\_integrators.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025/diff?src=pr&el=tree#diff-cGxhc21hcHkvc2ltdWxhdGlvbi9wYXJ0aWNsZV9pbnRlZ3JhdG9ycy5weQ==) | `33.33% <0.00%> (-66.67%)` | :arrow_down: |
| [plasmapy/simulation/particletracker.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025/diff?src=pr&el=tree#diff-cGxhc21hcHkvc2ltdWxhdGlvbi9wYXJ0aWNsZXRyYWNrZXIucHk=) | `36.20% <0.00%> (-62.07%)` | :arrow_down: |
| [plasmapy/plasma/sources/plasma3d.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025/diff?src=pr&el=tree#diff-cGxhc21hcHkvcGxhc21hL3NvdXJjZXMvcGxhc21hM2QucHk=) | `93.75% <0.00%> (-6.25%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=footer). Last update [4196ab7...3ba5a5c](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1025?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

StanczakDominik: Oh, and of course, forgot the most important thing!

Thank you! :)
StanczakDominik: ... and there's the missing changelog entry ;)
namurphy: And I got it in before midnight Eastern time, just as promised!  While listening to very loud music, a habit which is undoubtedly not the reason why my sleep schedule is drifting later.  I didn't learn about `yield from` until last month, it turns out.  Really helpful!

Still need to add a quick changelog entry.  I'll try not to let song lyrics slip in.

Patch:
diff --git a/changelog/1025.bugfix.rst b/changelog/1025.bugfix.rst
new file mode 100644
index 00000000..bef49d7b
--- /dev/null
+++ b/changelog/1025.bugfix.rst
@@ -0,0 +1,3 @@
+Fixed a bug that prevented nested iterations of a single
+`~plasmapy.particles.IonizationState` or
+`~plasmapy.particles.IonizationStateCollection` instance.
diff --git a/plasmapy/particles/ionization_state.py b/plasmapy/particles/ionization_state.py
index 03397194..82852d06 100644
--- a/plasmapy/particles/ionization_state.py
+++ b/plasmapy/particles/ionization_state.py
@@ -342,26 +342,7 @@ def __setitem__(self, key, value):
         )
 
     def __iter__(self):
-        """Initialize an instance prior to iteration."""
-        self._charge_index = 0
-        return self
-
-    def __next__(self):
-        """
-        Return a `~plasmapy.particles.IonicFraction` instance that contains
-        information about a particular ionization level.
-        """
-        if self._charge_index <= self.atomic_number:
-            result = IonicFraction(
-                ion=Particle(self.base_particle, Z=self._charge_index),
-                ionic_fraction=self.ionic_fractions[self._charge_index],
-                number_density=self.number_densities[self._charge_index],
-            )
-            self._charge_index += 1
-            return result
-        else:
-            del self._charge_index
-            raise StopIteration
+        yield from [self[i] for i in range(self.atomic_number + 1)]
 
     def __eq__(self, other):
         """
diff --git a/plasmapy/particles/ionization_state_collection.py b/plasmapy/particles/ionization_state_collection.py
index 32d1ed44..c374d423 100644
--- a/plasmapy/particles/ionization_state_collection.py
+++ b/plasmapy/particles/ionization_state_collection.py
@@ -335,28 +335,7 @@ def __setitem__(self, key, value):
         self._ionic_fractions[particle][:] = new_fractions[:]
 
     def __iter__(self):
-        """
-        Prepare an `~plasmapy.particles.IonizationStateCollection` instance for
-        iteration.
-        """
-        self._element_index = 0
-        return self
-
-    def __next__(self):
-        if self._element_index < len(self.base_particles):
-            particle = self.base_particles[self._element_index]
-            result = IonizationState(
-                particle,
-                self.ionic_fractions[particle],
-                T_e=self.T_e,
-                n_elem=np.sum(self.number_densities[particle]),
-                tol=self.tol,
-            )
-            self._element_index += 1
-            return result
-        else:
-            del self._element_index
-            raise StopIteration
+        yield from [self[key] for key in self.ionic_fractions.keys()]
 
     def __eq__(self, other):
 


Test patch:
diff --git a/plasmapy/particles/tests/test_ionization_collection.py b/plasmapy/particles/tests/test_ionization_collection.py
index d252307e..a69e4478 100644
--- a/plasmapy/particles/tests/test_ionization_collection.py
+++ b/plasmapy/particles/tests/test_ionization_collection.py
@@ -902,3 +902,15 @@ def test_number_density_assignment():
     instance = IonizationStateCollection(["H", "He"])
     number_densities = [2, 3, 5] * u.m ** -3
     instance["He"] = number_densities
+
+
+def test_iteration_with_nested_iterator():
+    ionization_states = IonizationStateCollection(["H", "He"])
+
+    i = 0
+    for ionization_state1 in ionization_states:
+        assert isinstance(ionization_state1, IonizationState)
+        for ionization_state2 in ionization_states:
+            assert isinstance(ionization_state2, IonizationState)
+            i += 1
+    assert i == 4
diff --git a/plasmapy/particles/tests/test_ionization_state.py b/plasmapy/particles/tests/test_ionization_state.py
index c907494b..6a652a6a 100644
--- a/plasmapy/particles/tests/test_ionization_state.py
+++ b/plasmapy/particles/tests/test_ionization_state.py
@@ -709,7 +709,6 @@ def test_kappa_isinf_when_not_set(self):
         assert np.isinf(self.instance.kappa)
 
 
-@pytest.mark.xfail
 def test_iteration_with_nested_iterator():
     hydrogen = IonizationState("p+", n_elem=1e20 * u.m ** -3, T_e=10 * u.eV)
 
@@ -717,6 +716,4 @@ def test_iteration_with_nested_iterator():
     for fraction in hydrogen:
         for fraction2 in hydrogen:
             i += 1
-            print(i, fraction, fraction2)
-            # should output 4 pairs...
     assert i == 4