Merge pull request #102 from compomics/fix/from-iter

PSMList: Use np.fromiter when using a string as accessor (e.g. `psm_list["peptidoform"]`)

Author: RalfG

psm_utils/psm_list.py

@@ -98,12 +98,7 @@ def __getitem__(self, item) -> PSM | list[PSM]:
             return PSMList(psm_list=self.psm_list[item])
         elif isinstance(item, str):
             # Return PSM property as array across full PSMList
-            try:
-                # Let NumPy coerce dtype (e.g., multidimensional arrays)
-                return np.array([psm[item] for psm in self.psm_list])
-            except ValueError:
-                # If dtype is not consistent, force dtype to be object
-                return np.array([psm[item] for psm in self.psm_list], dtype=object)
+            return np.fromiter([psm[item] for psm in self.psm_list], dtype=object, count=len(self))
         elif _is_iterable_of_bools(item):
             # Return new PSMList with items that were True
             return PSMList(psm_list=[self.psm_list[i] for i in np.flatnonzero(item)])

tests/test_psm_list.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from psm_utils import Peptidoform, PSM, PSMList
+from psm_utils import PSM, Peptidoform, PSMList
 
 sample_psm_list = [
     PSM(peptidoform="ACDK", spectrum_id=1, score=140.2),
@@ -42,7 +42,7 @@ def test___get_item__(self):
         # Multiple PSM properties as 2D array
         np.testing.assert_equal(
             psm_list[["spectrum_id", "score"]],
-            np.array([["1", 140.2], ["2", 132.9], ["3", 55.7]]),
+            np.array([["1", 140.2], ["2", 132.9], ["3", 55.7]], dtype=object),
         )
 
         # Index by multiple indices