__iter__ function for elements and fluid updated to pass mypy type check. test added

src/gaspype/_main.py

@@ -525,7 +525,9 @@ class fluid:
             return self.array_composition[..., [self.fs.species.index(k) if isinstance(k, str) else k for k in key]]
 
     def __iter__(self) -> Iterator[dict[str, float]]:
-        return iter({s: c for s, c in zip(self.fs.species, spa)} for spa in np.array(self.array_composition, ndmin=2))
+        assert len(self.shape) < 2, 'Cannot iterate over species with more than one dimension'
+        aec = self.array_composition.reshape(-1, len(self.fs.species))
+        return iter({s: c for s, c in zip(self.fs.species, aec[i, :])} for i in range(aec.shape[0]))
 
     def __repr__(self) -> str:
         if len(self.array_fractions.shape) == 1:
@@ -652,7 +654,9 @@ class elements:
             return self.array_elemental_composition[..., [self.fs.elements.index(k) if isinstance(k, str) else k for k in key]]
 
     def __iter__(self) -> Iterator[dict[str, float]]:
-        return iter({s: c for s, c in zip(self.fs.elements, spa)} for spa in np.array(self.array_elemental_composition, ndmin=2))
+        assert len(self.shape) < 2, 'Cannot iterate over elements with more than one dimension'
+        aec = self.array_elemental_composition.reshape(-1, len(self.fs.elements))
+        return iter({s: c for s, c in zip(self.fs.elements, aec[i, :])} for i in range(aec.shape[0]))
 
     def __repr__(self) -> str:
         if len(self.array_elemental_composition.shape) == 1:

tests/test_to_pandas.py

@@ -21,3 +21,13 @@ def test_elements():
 
     df = pd.DataFrame(list(gp.elements(fl * np.array([1, 2, 3, 4]))))
     assert df.shape == (4, 2)
+
+def test_iter():
+    fl = gp.fluid({'O2': 1, 'H2': 2, 'H2O': 3})
+
+    fl2 = fl * np.array([1, 2, 3, 4])
+    for i, f in enumerate(fl2):
+        if i == 1:
+            assert f == {'O2': np.float64(2.0), 'H2': np.float64(4.0), 'H2O': np.float64(6.0)}
+        if i == 3:
+            assert f == {'O2': np.float64(4.0), 'H2': np.float64(8.0), 'H2O': np.float64(12.0)}