Source code for CAT.attachment.remove_atoms_cm

"""A context manager for temporary removing a set of atoms from a molecule.

.. currentmodule:: CAT.attachment.remove_atoms_cm
.. autosummary::

.. autoclass:: RemoveAtoms


import sys
from typing import Iterable, Union, Sequence, cast
from contextlib import AbstractContextManager

from scm.plams import Molecule, Atom, Bond, MoleculeError

if sys.version_info >= (3, 8):
    # dict.__reversed__() was added in Python 3.8
    from builtins import dict as OrderedDict  # noqa: N812
    from collections import OrderedDict

__all__ = ['RemoveAtoms']

[docs]class RemoveAtoms(AbstractContextManager): """A context manager for temporary removing a set of atoms from a molecule. The *relative* ordering of the to-be removed atoms (and matching bonds), as specified in **atoms**, is preserved during the removal and reattachment process. Note that reattaching will (re-)append the removed atoms/bonds, a process which is thus likelly to affect the *absolute* ordering of atoms/bonds within the entire molecule. Examples -------- .. code:: python >>> from scm.plams import Molecule, Atom, from_smiles >>> mol: Molecule = from_smiles('CO') >>> atom1: Atom = mol[1] >>> atom2: Atom = mol[2] >>> atom_set = {atom1, atom2} >>> with RemoveAtoms(mol, atom_set): ... print(atom1 in mol, atom2 in mol) False False >>> print(atom1 in mol, atom2 in mol) True True Parameters ---------- mol : |plams.Molecule| A PLAMS molecule. See :attr:`RemoveAtoms.mol`. atoms : |plams.Atom| or |Iterable| [|plams.Atom|] A PLAMS atom or an iterable consisting of unique PLAMS atoms. All supplied atoms should belong to **mol**. See :attr:`RemoveAtoms.atoms`. Attributes ---------- mol : |plams.Molecule| A PLAMS molecule. atoms : |Sequence| [|plams.Atom|] A sequence of PLAMS atoms belonging to :attr:`RemoveAtoms.mol`. Setting a value will convert it into a sequence of atoms. _bonds : |OrderedDict| [|plams.Bond|, ``None``] A ordered dictionary of PLAMS bonds connected to one or more atoms in :attr:`RemoveAtoms.atoms`. All values are ``None``, the dictionary serving as an improvised ``OrderedSet``. Set to ``None`` until :meth:`RemoveAtoms.__enter__` is called. """ @property def atoms(self) -> Sequence[Atom]: """Get or set :attr:`RemoveAtoms.atoms`. Setting will convert the supplied value into a sequence of atoms.""" # noqa return self._atoms @atoms.setter def atoms(self, value: Union[Atom, Iterable[Atom]]) -> None: if isinstance(value, Atom): # It's an atom self._atoms = (value,) elif not hasattr(value, '__reversed__'): # It's an Iterator or Collection self._atoms = tuple(value) else: self._atoms = value # It's a Sequence (probably) def __init__(self, mol: Molecule, atoms: Union[Atom, Iterable[Atom]]) -> None: """Initialize a :class:`RemoveAtoms` instance.""" self.mol = mol self.atoms = cast(Sequence[Atom], atoms) self._bonds: 'OrderedDict[Bond, None]' = OrderedDict() # An improvised "OrderedSet" def __enter__(self) -> None: """Enter the context manager; delete all atoms in :class:`RemoveAtoms.atoms`.""" mol = self.mol bonds_set = self._bonds for atom in self.atoms: for bond in atom.bonds: bonds_set[bond] = None mol.delete_atom(atom) def __exit__(self, exc_type, exc_value, traceback) -> None: """Exit the context manager; reassign all atoms in :class:`RemoveAtoms.atoms`.""" mol = self.mol for atom in reversed(self.atoms): mol.add_atom(atom) for bond in reversed(self._bonds): try: mol.add_bond(bond) except MoleculeError: pass # One of the bonded atoms has been manually deleted by the user: skip it self._bonds = OrderedDict()