Fix calculate_mass with composition kwarg

levitsky · Feb 13, 2024 · 4f4a502 · 4f4a502
1 parent 74944e6
commit 4f4a502
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -21,7 +21,6 @@
    (`#138 <https://github.com/levitsky/pyteomics/pull/138>`_).
 
 
-
 4.6.3
 -----
 

diff --git a/pyteomics/mass/mass.py b/pyteomics/mass/mass.py
@@ -464,14 +464,13 @@ def mass(self, **kwargs):
         -------
         mass : float
         """
-        composition = self
         mass_data = kwargs.get('mass_data', nist_mass)
 
         # Calculate mass
         mass = 0.0
         average = kwargs.get('average', False)
 
-        for isotope_string, amount in composition.items():
+        for isotope_string, amount in self.items():
             element_name, isotope_num = _parse_isotope_string(isotope_string)
             # Calculate average mass if required and the isotope number is
             # not specified.
@@ -649,7 +648,10 @@ def calculate_mass(*args, **kwargs):
         if k in kwargs:
             mass_kw[k] = kwargs.pop(k)
     # Make a copy of `composition` keyword argument.
-    composition = (Composition(kwargs['composition']) if 'composition' in kwargs else Composition(*args, **kwargs))
+    if 'composition' in kwargs:
+        composition = Composition(kwargs.pop('composition'))
+    else:
+        composition = Composition(*args, **kwargs)
     kwargs.update(mass_kw)
     return composition.mass(**kwargs)
 

diff --git a/tests/test_mass.py b/tests/test_mass.py
@@ -148,6 +148,11 @@ def test_calculate_mass(self):
             mass.calculate_mass(parsed_sequence=['H-', 'X', 'Y', 'Z', '-OH'], aa_comp=self.aa_comp, mass_data=self.mass_data),
             sum(self.mass_data[atom][0][0] for atom in 'ABCDE'))
 
+        # Calculate mass by composition
+        self.assertEqual(
+            mass.calculate_mass(composition={'A': 1, 'B': 1, 'C': 1, 'D': 1, 'E': 1}, mass_data=self.mass_data),
+            sum(self.mass_data[atom][0][0] for atom in 'ABCDE'))
+
         # Calculate average mass by a formula.
         self.assertEqual(
             mass.calculate_mass(formula='ABCDE', average=True, mass_data=self.mass_data),
@@ -175,6 +180,10 @@ def test_calculate_mass(self):
                 mass.calculate_mass(formula='ABCDE', ion_type='M', charge=charge, mass_data=self.mass_data),
                 (mass.calculate_mass(formula='ABCDE', mass_data=self.mass_data) + self.mass_data['H+'][0][0] * charge) / charge)
 
+            self.assertAlmostEqual(
+                mass.calculate_mass(composition={'A': 1, 'B': 1, 'C': 1, 'D': 1, 'E': 1}, charge=charge, charge_carrier='BC+', mass_data=self.mass_data),
+                mass.calculate_mass(composition={'A': 1, 'B': 1 + charge, 'C': 1 + charge, 'D': 1, 'E': 1}, mass_data=self.mass_data) / charge)
+
             self.assertRaises(auxiliary.PyteomicsError, mass.calculate_mass, **{'formula': 'ABCDEH+%d' % charge,
                    'ion_type': 'M', 'charge': charge, 'mass_data': self.mass_data})