From 02cc6ef6d487f38d875242c49e9da44afcf95c76 Mon Sep 17 00:00:00 2001
From: RalfG <ralfg@hotmail.be>
Date: Tue, 19 Sep 2023 19:18:35 +0200
Subject: [PATCH] peptidoform: Add iter and len methods

---
 psm_utils/peptidoform.py  |  8 ++++++++
 tests/test_peptidoform.py | 26 ++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/psm_utils/peptidoform.py b/psm_utils/peptidoform.py
index 814a8df..fd9695d 100644
--- a/psm_utils/peptidoform.py
+++ b/psm_utils/peptidoform.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Iterable, List, Tuple, Union
+
 import numpy as np
 from pyteomics import mass, proforma
 
@@ -69,6 +71,12 @@ def __eq__(self, __o: Peptidoform) -> bool:
         except AttributeError:
             raise NotImplementedError("Object is not a Peptidoform")
 
+    def __iter__(self) -> Iterable[Tuple[str, Union[None, List[proforma.TagBase]]]]:
+        return self.parsed_sequence.__iter__()
+
+    def __len__(self) -> int:
+        return self.parsed_sequence.__len__()
+
     @property
     def proforma(self) -> str:
         """
diff --git a/tests/test_peptidoform.py b/tests/test_peptidoform.py
index d93002a..a9cb691 100644
--- a/tests/test_peptidoform.py
+++ b/tests/test_peptidoform.py
@@ -1,7 +1,33 @@
+from pyteomics import proforma
+
 from psm_utils.peptidoform import Peptidoform
 
 
 class TestPeptidoform:
+
+    def test__len__(self):
+        test_cases = [
+            ("ACDEFGHIK", 9),
+            ("[ac]-AC[cm]DEFGHIK", 9),
+            ("[ac]-AC[Carbamidomethyl]DEFGHIK", 9),
+            ("[Acetyl]-AC[cm]DEFGK", 7),
+            ("<[cm]@C>[Acetyl]-ACDK", 4),
+            ("<[Carbamidomethyl]@C>[ac]-ACDEFGHIK", 9),
+        ]
+
+        for test_case_in, expected_out in test_cases:
+            peptidoform = Peptidoform(test_case_in)
+            assert len(peptidoform) == expected_out
+
+    def test__iter__(self):
+        for aa, mods in Peptidoform("ACDEM[U:35]K"):
+            assert isinstance(aa, str)
+            if mods is not None:
+                assert isinstance(mods, list)
+                for mod in mods:
+                    assert isinstance(mod, proforma.TagBase)
+
+
     def test_rename_modifications(self):
         label_mapping = {"ac": "Acetyl", "cm": "Carbamidomethyl"}