From b020ad7eb0245d2e9aa76ff9930ccd82a4eb4764 Mon Sep 17 00:00:00 2001
From: Cajac102 <ca.jachmann@gmail.com>
Date: Fri, 9 Feb 2024 13:02:14 +0100
Subject: [PATCH 1/2] add tdf support with timsrust_pyo3

---
 ms2pip/spectrum_input.py | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/ms2pip/spectrum_input.py b/ms2pip/spectrum_input.py
index 190f527..0755c23 100644
--- a/ms2pip/spectrum_input.py
+++ b/ms2pip/spectrum_input.py
@@ -4,6 +4,7 @@
 from typing import Generator
 
 import numpy as np
+import timsrust_pyo3 as timsrust
 from pyteomics import mgf, mzml
 
 from ms2pip.exceptions import UnsupportedSpectrumFiletypeError
@@ -83,9 +84,31 @@ def read_mzml(spectrum_file: str) -> Generator[ObservedSpectrum, None, None]:
                 yield spectrum
 
 
+def read_tdf(spectrum_file: str) -> Generator[ObservedSpectrum, None, None]:
+    """
+    Read MS2 DDA spectra from .d folder.
+
+    Parameters
+    ----------
+    spectrum_file
+        Path to .d folder.
+
+    """
+    reader = timsrust.TimsReader(spectrum_file)
+    for spectrum in reader.read_all_spectra():
+        spectrum = ObservedSpectrum(
+            identifier=spectrum.index,
+            mz=np.asarray(spectrum.mz_values),
+            intensity=np.asarray(spectrum.intensities),
+            precursor_mz=spectrum.precursor.mz,
+            precursor_charge=spectrum.precursor.charge
+            )
+        yield spectrum
+
+
 def read_spectrum_file(spectrum_file: str) -> Generator[ObservedSpectrum, None, None]:
     """
-    Read MS2 spectra from MGF or mzML file; inferring the type from the filename extension.
+    Read MS2 spectra from MGF or mzML file or .d folder; inferring the type from the filename extension.
 
     Parameters
     ----------
@@ -100,5 +123,8 @@ def read_spectrum_file(spectrum_file: str) -> Generator[ObservedSpectrum, None,
     elif filetype == ".mgf":
         for spectrum in read_mgf(spectrum_file):
             yield spectrum
+    elif filetype == ".d":
+        for spectrum in read_tdf(spectrum_file):
+            yield spectrum
     else:
         raise UnsupportedSpectrumFiletypeError(filetype)

From e9bd5c0a732ef181922b70f9457b995e4cc5af58 Mon Sep 17 00:00:00 2001
From: Cajac102 <ca.jachmann@gmail.com>
Date: Fri, 9 Feb 2024 14:18:11 +0100
Subject: [PATCH 2/2] add timsrust_pyo3 as optional dependency

---
 ms2pip/spectrum_input.py | 8 +++++++-
 pyproject.toml           | 1 +
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/ms2pip/spectrum_input.py b/ms2pip/spectrum_input.py
index 0755c23..e951da2 100644
--- a/ms2pip/spectrum_input.py
+++ b/ms2pip/spectrum_input.py
@@ -4,7 +4,11 @@
 from typing import Generator
 
 import numpy as np
-import timsrust_pyo3 as timsrust
+try:
+    import timsrust_pyo3 as timsrust
+    _has_timsrust = True
+except ImportError:
+    _has_timsrust = False
 from pyteomics import mgf, mzml
 
 from ms2pip.exceptions import UnsupportedSpectrumFiletypeError
@@ -94,6 +98,8 @@ def read_tdf(spectrum_file: str) -> Generator[ObservedSpectrum, None, None]:
         Path to .d folder.
 
     """
+    if not _has_timsrust:
+        raise ImportError("Optional dependency timsrust_pyo3 required for .d spectrum file support.")
     reader = timsrust.TimsReader(spectrum_file)
     for spectrum in reader.read_all_spectra():
         spectrum = ObservedSpectrum(
diff --git a/pyproject.toml b/pyproject.toml
index d5907e8..1444ba1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,6 +50,7 @@ dependencies = [
 
 [project.optional-dependencies]
 plotting = ["matplotlib>=3.0", "spectrum-utils>=0.4"]
+tdf = ["timsrust_pyo3@git+https://github.com/jspaezp/timsrust_pyo3#egg=0f40c31"]
 dev = ["black", "isort>5", "pytest"]
 docs = [
     "sphinx",