Fix support for sqlalchemy v2, keeping backwards compatibility with v1.4

Author: RalfG

ms2pip/_utils/dlib.py

@@ -18,16 +18,18 @@
     TypeDecorator,
 )
 from sqlalchemy.dialects.sqlite import BLOB
+from sqlalchemy.engine import Connection
 
 DLIB_VERSION = "0.1.14"
 
 
 class CompressedArray(TypeDecorator):
-    """ Sqlite-like does not support arrays.
-        Let's use a custom type decorator.
+    """Sqlite-like does not support arrays.
+    Let's use a custom type decorator.
 
-        See http://docs.sqlalchemy.org/en/latest/core/types.html#sqlalchemy.types.TypeDecorator
+    See http://docs.sqlalchemy.org/en/latest/core/types.html#sqlalchemy.types.TypeDecorator
     """
+
     impl = BLOB
 
     def __init__(self, dtype, *args, **kwargs):
@@ -49,51 +51,55 @@ def copy(self):
 
 metadata = MetaData()
 
-big_float = numpy.dtype('>f4')
-big_double = numpy.dtype('>f8')
+big_float = numpy.dtype(">f4")
+big_double = numpy.dtype(">f8")
 
 Entry = Table(
-    'entries',
+    "entries",
     metadata,
-    Column('PrecursorMz', Float, nullable=False, index=True),
-    Column('PrecursorCharge', Integer, nullable=False),
-    Column('PeptideModSeq', String, nullable=False),
-    Column('PeptideSeq', String, nullable=False, index=True),
-    Column('Copies', Integer, nullable=False),
-    Column('RTInSeconds', Float, nullable=False),
-    Column('Score', Float, nullable=False),
-    Column('MassEncodedLength', Integer, nullable=False),
-    Column('MassArray', CompressedArray(big_double), nullable=False),
-    Column('IntensityEncodedLength', Integer, nullable=False),
-    Column('IntensityArray', CompressedArray(big_float), nullable=False),
-    Column('CorrelationEncodedLength', Integer, nullable=True),
-    Column('CorrelationArray', CompressedArray(big_float), nullable=True),
-    Column('RTInSecondsStart', Float, nullable=True),
-    Column('RTInSecondsStop', Float, nullable=True),
-    Column('MedianChromatogramEncodedLength', Integer, nullable=True),
-    Column('MedianChromatogramArray', CompressedArray(big_float), nullable=True),
-    Column('SourceFile', String, nullable=False),
+    Column("PrecursorMz", Float, nullable=False, index=True),
+    Column("PrecursorCharge", Integer, nullable=False),
+    Column("PeptideModSeq", String, nullable=False),
+    Column("PeptideSeq", String, nullable=False, index=True),
+    Column("Copies", Integer, nullable=False),
+    Column("RTInSeconds", Float, nullable=False),
+    Column("Score", Float, nullable=False),
+    Column("MassEncodedLength", Integer, nullable=False),
+    Column("MassArray", CompressedArray(big_double), nullable=False),
+    Column("IntensityEncodedLength", Integer, nullable=False),
+    Column("IntensityArray", CompressedArray(big_float), nullable=False),
+    Column("CorrelationEncodedLength", Integer, nullable=True),
+    Column("CorrelationArray", CompressedArray(big_float), nullable=True),
+    Column("RTInSecondsStart", Float, nullable=True),
+    Column("RTInSecondsStop", Float, nullable=True),
+    Column("MedianChromatogramEncodedLength", Integer, nullable=True),
+    Column("MedianChromatogramArray", CompressedArray(big_float), nullable=True),
+    Column("SourceFile", String, nullable=False),
 )
 
-Index('ix_entries_PeptideModSeq_PrecursorCharge_SourceFile', Entry.c.PeptideModSeq, Entry.c.PrecursorCharge, Entry.c.SourceFile)
+Index(
+    "ix_entries_PeptideModSeq_PrecursorCharge_SourceFile",
+    Entry.c.PeptideModSeq,
+    Entry.c.PrecursorCharge,
+    Entry.c.SourceFile,
+)
 
 PeptideToProtein = Table(
-    'peptidetoprotein',
+    "peptidetoprotein",
     metadata,
-    Column('PeptideSeq', String, nullable=False, index=True),
-    Column('isDecoy', Boolean, nullable=True),
-    Column('ProteinAccession', String, nullable=False, index=True),
+    Column("PeptideSeq", String, nullable=False, index=True),
+    Column("isDecoy", Boolean, nullable=True),
+    Column("ProteinAccession", String, nullable=False, index=True),
 )
 
 Metadata = Table(
-    'metadata',
+    "metadata",
     metadata,
-    Column('Key', String, nullable=False, index=True),
-    Column('Value', String, nullable=False),
+    Column("Key", String, nullable=False, index=True),
+    Column("Value", String, nullable=False),
 )
 
 
-def open_sqlite(filename: Union[str, Path]) -> sqlalchemy.engine.Connection:
+def open_sqlite(filename: Union[str, Path]) -> Connection:
     engine = sqlalchemy.create_engine(f"sqlite:///{filename}")
-    metadata.bind = engine
     return engine.connect()

ms2pip/spectrum_output.py

@@ -53,7 +53,8 @@
 import numpy as np
 from psm_utils import PSM, Peptidoform
 from pyteomics import proforma
-from sqlalchemy import engine, select
+from sqlalchemy import select
+from sqlalchemy.engine import Connection
 
 from ms2pip._utils import dlib
 from ms2pip.result import ProcessingResult
@@ -653,9 +654,9 @@ def open(self):
     def write(self, processing_results: List[ProcessingResult]):
         """Write MS2PIP predictions to a DLIB SQLite file."""
         connection = self._file_object
-        dlib.metadata.create_all()
+        dlib.metadata.create_all(connection.engine)
         self._write_metadata(connection)
-        self._write_entries(processing_results, connection, self.filename)
+        self._write_entries(processing_results, connection, str(self.filename))
         self._write_peptide_to_protein(processing_results, connection)
 
     def _write_result(self, result: ProcessingResult): ...
@@ -682,11 +683,11 @@ def _format_modified_sequence(peptidoform: Peptidoform) -> str:
         )
 
     @staticmethod
-    def _write_metadata(connection: engine.Connection):
+    def _write_metadata(connection: Connection):
         """Write metadata to DLIB SQLite file."""
         with connection.begin():
             version = connection.execute(
-                select([dlib.Metadata.c.Value]).where(dlib.Metadata.c.Key == "version")
+                select(dlib.Metadata.c.Value).where(dlib.Metadata.c.Key == "version")
             ).scalar()
             if version is None:
                 connection.execute(
@@ -699,7 +700,7 @@ def _write_metadata(connection: engine.Connection):
     @staticmethod
     def _write_entries(
         processing_results: List[ProcessingResult],
-        connection: engine.Connection,
+        connection: Connection,
         output_filename: str,
     ):
         """Write spectra to DLIB SQLite file."""
@@ -730,7 +731,7 @@ def _write_entries(
                 )
 
     @staticmethod
-    def _write_peptide_to_protein(results: List[ProcessingResult], connection: engine.Connection):
+    def _write_peptide_to_protein(results: List[ProcessingResult], connection: Connection):
         """Write peptide-to-protein mappings to DLIB SQLite file."""
         peptide_to_proteins = {
             (result.psm.peptidoform.sequence, protein)
@@ -743,7 +744,7 @@ def _write_peptide_to_protein(results: List[ProcessingResult], connection: engin
             sql_peptide_to_proteins = set()
             proteins = {protein for _, protein in peptide_to_proteins}
             for peptide_to_protein in connection.execute(
-                dlib.PeptideToProtein.select().where(
+                select(dlib.PeptideToProtein).where(
                     dlib.PeptideToProtein.c.ProteinAccession.in_(proteins)
                 )
             ):

pyproject.toml

@@ -37,8 +37,7 @@ dependencies = [
     "pandas>=1,<3",
     "pyarrow",
     "pyteomics>=3.5,<5",
-    "tomlkit>=0.5,<1",
-    "sqlalchemy>=1.3,<2",
+    "sqlalchemy>=1.4,<3",
     "click>=7,<9",
     "xgboost>=1.3,<3",
     "lxml>=4",