Skip to content

Commit

Permalink
fix Composition.__contains__() by handling Species and DummySpecies c…
Browse files Browse the repository at this point in the history
…orrectly
  • Loading branch information
janosh committed Jul 25, 2023
1 parent 3783963 commit d86e774
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions pymatgen/core/composition.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,10 @@ def __iter__(self) -> Iterator[Species | Element | DummySpecies]:
def __contains__(self, key) -> bool:
try:
sp = get_el_sp(key)
return sp in self._data
if isinstance(sp, Species):
return sp in self._data
# Element or str
return any(sp.symbol == s.symbol for s in self._data)
except ValueError as exc:
raise TypeError(f"Invalid {key=} for Composition") from exc

Expand All @@ -160,9 +163,9 @@ def __eq__(self, other: object) -> bool:
if not isinstance(other, (Composition, dict)):
return NotImplemented

# elements with amounts < Composition.amount_tolerance don't show up
# in the el_map, so checking len enables us to only check one
# composition's elements
# elements with amounts < Composition.amount_tolerance don't show up
# in the el_map, so checking len enables us to only check one
# composition's elements
if len(self) != len(other):
return False

Expand Down

0 comments on commit d86e774

Please sign in to comment.