Handle mixed types in ResidueId comparison (#235)

* Update residue.py * chore: formatting * add tests * chore: add entry to changelog --------- Co-authored-by: Cédric Bouysset <[email protected]>
chemosim-lab · Dec 23, 2024 · 742203b · 742203b
1 parent 597083d
commit 742203b
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -35,6 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Sorting `ResidueId` objects where a subset of residues had no chain was raising a
+  `TypeError`, it will now put cases without a chain first (PR #235 by @amorehead).
 - `display_residues` was sanitizing each residue while preparing them for display, which
   could make debugging faulty molecules difficult. This is now disabled.
 - Deprecation warnings

diff --git a/prolif/residue.py b/prolif/residue.py
@@ -66,7 +66,15 @@ def __eq__(self, other: "ResidueId"):
         )
 
     def __lt__(self, other: "ResidueId"):
-        return (self.chain, self.number) < (other.chain, other.number)
+        def chain_key(chain):
+            # Handles the case where the two chains are of different types
+            # e.g., None from WAT123, and str from ALA42.A
+            return (chain is not None, chain)
+
+        return (chain_key(self.chain), self.number) < (
+            chain_key(other.chain),
+            other.number,
+        )
 
     @classmethod
     def from_atom(cls, atom):

diff --git a/tests/test_residues.py b/tests/test_residues.py
@@ -124,6 +124,8 @@ def test_eq(self):
             ("ALA1.A", "ALA1.B"),
             ("ALA2.A", "ALA3.A"),
             ("ALA4.A", "ALA1.B"),
+            ("ALA1", "ALA2"),
+            ("ALA1", "ALA1.A"),
         ],
     )
     def test_lt(self, res1, res2):