diff --git a/openff/toolkit/topology/_mm_molecule.py b/openff/toolkit/topology/_mm_molecule.py index 42cd31f65..382102ef4 100644 --- a/openff/toolkit/topology/_mm_molecule.py +++ b/openff/toolkit/topology/_mm_molecule.py @@ -82,13 +82,13 @@ def n_bonds(self) -> int: def n_conformers(self) -> int: return 0 if self._conformers is None else len(self._conformers) - def atom(self, index): + def atom(self, index: int) -> "_SimpleAtom": return self.atoms[index] def atom_index(self, atom) -> int: return self.atoms.index(atom) - def bond(self, index): + def bond(self, index: int) -> "_SimpleBond": return self.bonds[index] def get_bond_between(self, atom1_index, atom2_index): @@ -340,9 +340,7 @@ def to_dict(self) -> dict: molecule_dict["conformers"] = None else: molecule_dict["conformers"] = [] - molecule_dict["conformers_unit"] = ( - "angstrom" # Have this defined as a class variable? - ) + molecule_dict["conformers_unit"] = "angstrom" # Have this defined as a class variable? for conf in self._conformers: conf_unitless = conf.m_as(unit.angstrom) conf_serialized, conf_shape = serialize_numpy(conf_unitless) @@ -368,9 +366,7 @@ def from_dict(cls, molecule_dict): for bond_dict in bond_dicts: atom1_index = bond_dict["atom1_index"] atom2_index = bond_dict["atom2_index"] - molecule.add_bond( - atom1=molecule.atom(atom1_index), atom2=molecule.atom(atom2_index) - ) + molecule.add_bond(atom1=molecule.atom(atom1_index), atom2=molecule.atom(atom2_index)) conformers = molecule_dict.pop("conformers") if conformers is None: @@ -438,9 +434,7 @@ def to_molecule(self) -> NoReturn: "an OpenFF Molecule with sufficiently specified chemistry." ) - def is_isomorphic_with( - self, other: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"], **kwargs - ) -> bool: + def is_isomorphic_with(self, other: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"], **kwargs) -> bool: """ Check for pseudo-isomorphism. @@ -504,9 +498,7 @@ def node_match_func(node1, node2): if return_atom_map: topology_atom_map = matcher.mapping - return True, { - key: topology_atom_map[key] for key in sorted(topology_atom_map) - } + return True, {key: topology_atom_map[key] for key in sorted(topology_atom_map)} else: return True, None @@ -534,9 +526,7 @@ def __getattr__(self, name: str) -> list["HierarchyElement"]: try: return self.__dict__["_hierarchy_schemes"][name].hierarchy_elements except KeyError: - raise AttributeError( - f"'{self.__class__.__name__}' object has no attribute {name!r}" - ) + raise AttributeError(f"'{self.__class__.__name__}' object has no attribute {name!r}") def __deepcopy__(self, memo): return self.__class__.from_dict(self.to_dict()) @@ -585,10 +575,7 @@ def atomic_number(self, value): if not isinstance(value, int): raise ValueError("atomic_number must be an integer") if value < 0: - raise ValueError( - "atomic_number must be non-negative. An atomic number " - "of 0 is acceptable." - ) + raise ValueError("atomic_number must be non-negative. An atomic number of 0 is acceptable.") self._atomic_number = value @property diff --git a/openff/toolkit/topology/molecule.py b/openff/toolkit/topology/molecule.py index 64f766f2f..f826fb64f 100644 --- a/openff/toolkit/topology/molecule.py +++ b/openff/toolkit/topology/molecule.py @@ -30,7 +30,7 @@ import pathlib import warnings from collections import UserDict, defaultdict -from collections.abc import Generator, Iterable, Sequence +from collections.abc import Generator, Iterable, Mapping, MutableMapping, Sequence from copy import deepcopy from functools import cmp_to_key from typing import ( @@ -91,6 +91,7 @@ import IPython.display import networkx as nx import nglview + from openmm.unit import Quantity as OMMQuantity from rdkit.Chem import Mol as RDMol from openff.toolkit.topology._mm_molecule import _SimpleAtom, _SimpleMolecule @@ -148,9 +149,7 @@ def molecule(self, molecule: "FrozenMolecule"): Set the particle's molecule pointer. Note that this will only work if the particle currently doesn't have a molecule """ - assert ( - self._molecule is None - ), f"{type(self).__name__} already has an associated molecule" + assert self._molecule is None, f"{type(self).__name__} already has an associated molecule" self._molecule = molecule @property @@ -158,7 +157,7 @@ def molecule_particle_index(self) -> int: """ Returns the index of this particle in its molecule """ - return self._molecule.atoms.index(self) + return self._molecule.atoms.index(self) # type:ignore @property def name(self) -> str: @@ -193,13 +192,10 @@ def __init__(self, *args, **kwargs): def __setitem__(self, key, value): if not isinstance(key, str): - raise InvalidAtomMetadataError( - f"Attempted to set atom metadata with a non-string key. (key: {key}" - ) + raise InvalidAtomMetadataError(f"Attempted to set atom metadata with a non-string key. (key: {key}") if not isinstance(value, (str, int)): raise InvalidAtomMetadataError( - f"Attempted to set atom metadata with a non-string or integer " - f"value. (value: {value})" + f"Attempted to set atom metadata with a non-string or integer value. (value: {value})" ) super().__setitem__(key, value) @@ -225,8 +221,8 @@ def __init__( is_aromatic: bool, name: Optional[str] = None, molecule=None, - stereochemistry: Optional[str] = None, - metadata: Optional[dict[str, Union[int, str]]] = None, + stereochemistry: Literal["R", "S", None] = None, + metadata: Mapping[str, int | str] | None = None, ): """ Create an immutable Atom object. @@ -271,7 +267,7 @@ def __init__( # Use the setter here, since it will handle either ints or Quantities # and it is designed to quickly process ints - self.formal_charge = formal_charge + self.formal_charge = formal_charge # type: ignore[assignment] self._is_aromatic = is_aromatic self._stereochemistry = stereochemistry if name is None: @@ -311,9 +307,13 @@ def to_dict(self) -> dict[str, Union[None, str, int, bool, dict[Any, Any]]]: """ # TODO: Should this be implicit in the atom ordering when saved? # atom_dict['molecule_atom_index'] = self._molecule_atom_index + + # trust that the unit is e + formal_charge: int = self._formal_charge.m # type: ignore + return { "atomic_number": self._atomic_number, - "formal_charge": self._formal_charge.m, # Trust that the unit is e + "formal_charge": formal_charge, "is_aromatic": self._is_aromatic, "stereochemistry": self._stereochemistry, "name": self._name, @@ -330,21 +330,21 @@ def from_dict(cls: type[A], atom_dict: dict) -> A: return cls(**atom_dict) @property - def metadata(self): + def metadata(self) -> MutableMapping[str, int | str]: """ The atom's metadata dictionary """ return self._metadata @property - def formal_charge(self): + def formal_charge(self) -> Quantity: """ The atom's formal charge """ return self._formal_charge @formal_charge.setter - def formal_charge(self, other): + def formal_charge(self, other: "int | Quantity | OMMQuantity"): """ Set the atom's formal charge. Accepts either ints or unit-wrapped ints with units of charge. """ @@ -355,16 +355,13 @@ def formal_charge(self, other): if other.units in _CHARGE_UNITS: self._formal_charge = other else: - raise IncompatibleUnitError( - f"Cannot set formal charge with a quantity with units {other.units}" - ) + raise IncompatibleUnitError(f"Cannot set formal charge with a quantity with units {other.units}") elif hasattr(other, "unit"): from openmm import unit as openmm_unit if not isinstance(other, openmm_unit.Quantity): raise IncompatibleUnitError( - "Unsupported type passed to formal_charge setter. " - f"Found object of type {type(other)}." + f"Unsupported type passed to formal_charge setter. Found object of type {type(other)}." ) from openff.units.openmm import from_openmm @@ -373,9 +370,7 @@ def formal_charge(self, other): if converted.units in _CHARGE_UNITS: self._formal_charge = converted else: - raise IncompatibleUnitError( - f"Cannot set formal charge with a quantity with units {converted.units}" - ) + raise IncompatibleUnitError(f"Cannot set formal charge with a quantity with units {converted.units}") else: raise ValueError @@ -431,19 +426,19 @@ def is_aromatic(self): return self._is_aromatic @property - def stereochemistry(self): + def stereochemistry(self) -> Literal["R", "S", None]: """ The atom's stereochemistry (if defined, otherwise None) """ return self._stereochemistry @stereochemistry.setter - def stereochemistry(self, value: Literal["CW", "CCW", None]): + def stereochemistry(self, value: Literal["R", "S", None]): """Set the atoms stereochemistry Parameters ---------- value - The stereochemistry around this atom, allowed values are "CW", "CCW", or None, + The stereochemistry around this atom, allowed values are "R", "S", or None, """ # if (value != 'CW') and (value != 'CCW') and not(value is None): @@ -499,9 +494,7 @@ def name(self, other: str): The new name for this atom """ if type(other) is not str: - raise ValueError( - f"In setting atom name. Expected str, received {other} (type {type(other)})." - ) + raise ValueError(f"In setting atom name. Expected str, received {other} (type {type(other)}).") self._name = other @property @@ -761,11 +754,11 @@ def atoms(self): return (self._atom1, self._atom2) @property - def bond_order(self): + def bond_order(self) -> int: return self._bond_order @bond_order.setter - def bond_order(self, value): + def bond_order(self, value: int): if isinstance(value, int): self._bond_order = value else: @@ -786,7 +779,7 @@ def fractional_bond_order(self, value): self._fractional_bond_order = value @property - def stereochemistry(self): + def stereochemistry(self) -> Literal["E", "Z", None]: return self._stereochemistry @property @@ -805,9 +798,7 @@ def molecule(self, value): # TODO: This is an impossible state (the constructor requires that atom1 and atom2 # are in a molecule, the same molecule, and sets that as self._molecule). # Should we remove this? - assert ( - self._molecule is None - ), "Bond.molecule is already set and can only be set once" + assert self._molecule is None, "Bond.molecule is already set and can only be set once" self._molecule = value @property @@ -852,9 +843,7 @@ def __repr__(self): return f"Bond(atom1 index={self.atom1_index}, atom2 index={self.atom2_index})" def __str__(self): - return ( - f"" - ) + return f"" # TODO: How do we automatically trigger invalidation of cached properties if an ``Atom`` or ``Bond`` is modified, @@ -1062,11 +1051,7 @@ def __init__( loaded = True # TODO: Make this compatible with file-like objects (I couldn't figure out how to make an oemolistream # from a fileIO object) - if ( - isinstance(other, (str, pathlib.Path)) - or (hasattr(other, "read") - and not loaded) - ): + if isinstance(other, (str, pathlib.Path)) or (hasattr(other, "read") and not loaded): try: mol = Molecule.from_file( other, @@ -1075,9 +1060,7 @@ def __init__( allow_undefined_stereo=allow_undefined_stereo, ) # returns a list only if multiple molecules are found if type(mol) is list: - raise ValueError( - "Specified file or file-like object must contain exactly one molecule" - ) + raise ValueError("Specified file or file-like object must contain exactly one molecule") except ValueError as e: value_errors.append(e) else: @@ -1088,9 +1071,7 @@ def __init__( # errors from the different loading attempts if not loaded: - msg = ( - f"Cannot construct openff.toolkit.topology.Molecule from {other}\n" - ) + msg = f"Cannot construct openff.toolkit.topology.Molecule from {other}\n" for value_error in value_errors: msg += str(value_error) raise ValueError(msg) @@ -1207,6 +1188,7 @@ def to_dict(self) -> dict: list[str], list[bytes], list[HierarchyElement], + list[dict[str, Any]], ], ] = dict() molecule_dict["name"] = self._name @@ -1223,19 +1205,14 @@ def to_dict(self) -> dict: molecule_dict["conformers"] = None else: molecule_dict["conformers_unit"] = "angstrom" - molecule_dict["conformers"] = [ - serialize_numpy(conf.m_as(unit.angstrom))[0] - for conf in self._conformers - ] + molecule_dict["conformers"] = [serialize_numpy(conf.m_as(unit.angstrom))[0] for conf in self._conformers] if self._partial_charges is None: molecule_dict["partial_charges"] = None molecule_dict["partial_charge_unit"] = None else: - molecule_dict["partial_charges"], _ = serialize_numpy( - self._partial_charges.m_as(unit.elementary_charge) - ) + molecule_dict["partial_charges"], _ = serialize_numpy(self._partial_charges.m_as(unit.elementary_charge)) molecule_dict["partial_charge_unit"] = "elementary_charge" molecule_dict["hierarchy_schemes"] = dict() @@ -1336,9 +1313,7 @@ def _initialize_from_dict(self, molecule_dict: dict): self._properties = deepcopy(molecule_dict["properties"]) - for iter_name, hierarchy_scheme_dict in molecule_dict[ - "hierarchy_schemes" - ].items(): + for iter_name, hierarchy_scheme_dict in molecule_dict["hierarchy_schemes"].items(): # It's important that we do NOT call `add_hierarchy_scheme` here, since we # need to deserialize these HierarchyElements exactly as they were serialized, # even if that conflicts with the current values in atom metadata. @@ -1350,9 +1325,7 @@ def _initialize_from_dict(self, molecule_dict: dict): self._hierarchy_schemes[iter_name] = new_hier_scheme for element_dict in hierarchy_scheme_dict["hierarchy_elements"]: - new_hier_scheme.add_hierarchy_element( - tuple(element_dict["identifier"]), element_dict["atom_indices"] - ) + new_hier_scheme.add_hierarchy_element(tuple(element_dict["identifier"]), element_dict["atom_indices"]) def __repr__(self): """Return a summary of this molecule; SMILES if valid, Hill formula if not.""" @@ -1368,9 +1341,9 @@ def _initialize(self): """ Clear the contents of the current molecule. """ - self._name = "" - self._atoms = list() - self._bonds = list() # list of bonds between Atom objects + self._name: str = "" + self._atoms: list[Atom] = list() + self._bonds: list[Bond] = list() # list of bonds between Atom objects self._properties = {} # Attached properties to be preserved # self._cached_properties = None # Cached properties (such as partial charges) can be recomputed as needed self._partial_charges = None @@ -1467,9 +1440,7 @@ def _add_residue_hierarchy_scheme(self, overwrite_existing: bool = True): if "residues" in self._hierarchy_schemes.keys(): self.delete_hierarchy_scheme("residues") - self.add_hierarchy_scheme( - ("chain_id", "residue_number", "insertion_code", "residue_name"), "residues" - ) + self.add_hierarchy_scheme(("chain_id", "residue_number", "insertion_code", "residue_name"), "residues") def add_hierarchy_scheme( self, @@ -1618,9 +1589,7 @@ def __getattr__(self, name: str) -> list["HierarchyElement"]: try: return self.__dict__["_hierarchy_schemes"][name].hierarchy_elements except KeyError: - raise AttributeError( - f"'{self.__class__.__name__}' object has no attribute {name!r}" - ) + raise AttributeError(f"'{self.__class__.__name__}' object has no attribute {name!r}") def __dir__(self): """Add the hierarchy scheme iterator names to dir""" @@ -1682,12 +1651,7 @@ def to_smiles( # Get a string representation of the function containing the toolkit name so we can check # if a SMILES was already cached for this molecule. This will return, for example # "RDKitToolkitWrapper.to_smiles" - smiles_hash = ( - to_smiles_method.__qualname__ - + str(isomeric) - + str(explicit_hydrogens) - + str(mapped) - ) + smiles_hash = to_smiles_method.__qualname__ + str(isomeric) + str(explicit_hydrogens) + str(mapped) smiles_hash += str(self._properties.get("atom_map", None)) # Check to see if a SMILES for this molecule was already cached using this method if smiles_hash in self._cached_smiles: @@ -1792,9 +1756,7 @@ def to_inchi( """ if isinstance(toolkit_registry, ToolkitRegistry): - inchi = toolkit_registry.call( - "to_inchi", self, fixed_hydrogens=fixed_hydrogens - ) + inchi = toolkit_registry.call("to_inchi", self, fixed_hydrogens=fixed_hydrogens) elif isinstance(toolkit_registry, ToolkitWrapper): toolkit = toolkit_registry inchi = toolkit.to_inchi(self, fixed_hydrogens=fixed_hydrogens) # type: ignore[attr-defined] @@ -1839,9 +1801,7 @@ def to_inchikey( """ if isinstance(toolkit_registry, ToolkitRegistry): - inchi_key = toolkit_registry.call( - "to_inchikey", self, fixed_hydrogens=fixed_hydrogens - ) + inchi_key = toolkit_registry.call("to_inchikey", self, fixed_hydrogens=fixed_hydrogens) elif isinstance(toolkit_registry, ToolkitWrapper): toolkit = toolkit_registry inchi_key = toolkit.to_inchikey(self, fixed_hydrogens=fixed_hydrogens) # type: ignore[attr-defined] @@ -1978,8 +1938,8 @@ def _is_exactly_the_same_as(self, other): @staticmethod def are_isomorphic( - mol1: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"], - mol2: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"], + mol1: "FrozenMolecule | _SimpleMolecule | nx.Graph[int]", + mol2: "FrozenMolecule | _SimpleMolecule | nx.Graph[int]", return_atom_map: bool = False, aromatic_matching: bool = True, formal_charge_matching: bool = True, @@ -1988,7 +1948,7 @@ def are_isomorphic( bond_stereochemistry_matching: bool = True, strip_pyrimidal_n_atom_stereo: bool = True, toolkit_registry: TKR = GLOBAL_TOOLKIT_REGISTRY, - ) -> tuple[bool, Optional[dict[int, int]]]: + ) -> tuple[bool, None | dict[int, int]]: """ Determine if ``mol1`` is isomorphic to ``mol2``. @@ -2093,8 +2053,7 @@ def _object_to_n_atoms(obj): return obj.number_of_nodes() else: raise TypeError( - "are_isomorphic accepts a NetworkX Graph or OpenFF " - + f"(Frozen)Molecule, not {type(obj)}" + "are_isomorphic accepts a NetworkX Graph or OpenFF " + f"(Frozen)Molecule, not {type(obj)}" ) # Quick number of atoms check. Important for large molecules @@ -2102,9 +2061,7 @@ def _object_to_n_atoms(obj): return False, None # If the number of atoms match, check the Hill formula - if Molecule._object_to_hill_formula(mol1) != Molecule._object_to_hill_formula( - mol2 - ): + if Molecule._object_to_hill_formula(mol1) != Molecule._object_to_hill_formula(mol2): return False, None # Do a quick check to see whether the inputs are totally identical (including being in the same atom order) @@ -2135,9 +2092,7 @@ def edge_match_func(x, y): # if the bond is aromatic. This way we avoid missing a match only # if the alternate bond orders 1 and 2 are assigned differently. if aromatic_matching and bond_order_matching: - is_equal = (x["is_aromatic"] == y["is_aromatic"]) or ( - x["bond_order"] == y["bond_order"] - ) + is_equal = (x["is_aromatic"] == y["is_aromatic"]) or (x["bond_order"] == y["bond_order"]) elif aromatic_matching: is_equal = x["is_aromatic"] == y["is_aromatic"] elif bond_order_matching: @@ -2166,9 +2121,7 @@ def to_networkx(data: Union[FrozenMolecule, nx.Graph]) -> nx.Graph: if strip_pyrimidal_n_atom_stereo: # Make a copy of the molecule so we don't modify the original data = deepcopy(data) - data.strip_atom_stereochemistry( - SMARTS, toolkit_registry=toolkit_registry - ) + data.strip_atom_stereochemistry(SMARTS, toolkit_registry=toolkit_registry) return data.to_networkx() elif isinstance(data, nx.Graph): @@ -2186,9 +2139,7 @@ def to_networkx(data: Union[FrozenMolecule, nx.Graph]) -> nx.Graph: from networkx.algorithms.isomorphism import GraphMatcher - GM = GraphMatcher( - mol1_netx, mol2_netx, node_match=node_match_func, edge_match=edge_match_func - ) + GM = GraphMatcher(mol1_netx, mol2_netx, node_match=node_match_func, edge_match=edge_match_func) isomorphic = GM.is_isomorphic() if isomorphic and return_atom_map: @@ -2206,8 +2157,14 @@ def to_networkx(data: Union[FrozenMolecule, nx.Graph]) -> nx.Graph: def is_isomorphic_with( self, - other: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"], - **kwargs, + other: "FrozenMolecule | _SimpleMolecule | nx.Graph[int]", + aromatic_matching: bool = True, + formal_charge_matching: bool = True, + bond_order_matching: bool = True, + atom_stereochemistry_matching: bool = True, + bond_stereochemistry_matching: bool = True, + strip_pyrimidal_n_atom_stereo: bool = True, + toolkit_registry: TKR = GLOBAL_TOOLKIT_REGISTRY, ) -> bool: """ Check if the molecule is isomorphic with the other molecule which can be an openff.toolkit.topology.Molecule @@ -2254,29 +2211,23 @@ def is_isomorphic_with( self, other, return_atom_map=False, - aromatic_matching=kwargs.get("aromatic_matching", True), - formal_charge_matching=kwargs.get("formal_charge_matching", True), - bond_order_matching=kwargs.get("bond_order_matching", True), - atom_stereochemistry_matching=kwargs.get( - "atom_stereochemistry_matching", True - ), - bond_stereochemistry_matching=kwargs.get( - "bond_stereochemistry_matching", True - ), - strip_pyrimidal_n_atom_stereo=kwargs.get( - "strip_pyrimidal_n_atom_stereo", True - ), - toolkit_registry=kwargs.get("toolkit_registry", GLOBAL_TOOLKIT_REGISTRY), + aromatic_matching=aromatic_matching, + formal_charge_matching=formal_charge_matching, + bond_order_matching=bond_order_matching, + atom_stereochemistry_matching=atom_stereochemistry_matching, + bond_stereochemistry_matching=bond_stereochemistry_matching, + strip_pyrimidal_n_atom_stereo=strip_pyrimidal_n_atom_stereo, + toolkit_registry=toolkit_registry, )[0] def generate_conformers( self, toolkit_registry: TKR = GLOBAL_TOOLKIT_REGISTRY, n_conformers: int = 10, - rms_cutoff: Optional[Quantity] = None, + rms_cutoff: Quantity | None = None, clear_existing: bool = True, make_carboxylic_acids_cis: bool = True, - ): + ) -> None: """ Generate conformers for this molecule using an underlying toolkit. @@ -2344,9 +2295,7 @@ def generate_conformers( f"Got {type(toolkit_registry)}" ) - def _make_carboxylic_acids_cis( - self, toolkit_registry: TKR = GLOBAL_TOOLKIT_REGISTRY - ): + def _make_carboxylic_acids_cis(self, toolkit_registry: TKR = GLOBAL_TOOLKIT_REGISTRY): """ Rotate dihedral angle of any conformers with trans COOH groups so they are cis @@ -2393,9 +2342,7 @@ def _make_carboxylic_acids_cis( conformers = np.asarray([q.m_as(unit.angstrom) for q in self._conformers]) # Scan the molecule for carboxylic acids - cooh_indices = self.chemical_environment_matches( - "[C:2]([O:3][H:4])=[O:1]", toolkit_registry=toolkit_registry - ) + cooh_indices = self.chemical_environment_matches("[C:2]([O:3][H:4])=[O:1]", toolkit_registry=toolkit_registry) n_conformers, n_cooh_groups = len(conformers), len(cooh_indices) # Exit early if there are no carboxylic acids if not n_cooh_groups: @@ -2450,9 +2397,7 @@ def dihedral(a): dihedrals.shape = (n_conformers, n_cooh_groups, 1, 1) # Get indices of trans COOH groups - trans_indices = np.logical_not( - np.logical_and((-np.pi / 2) < dihedrals, dihedrals < (np.pi / 2)) - ) + trans_indices = np.logical_not(np.logical_and((-np.pi / 2) < dihedrals, dihedrals < (np.pi / 2))) # Expand array so it can be used to index cooh_xyz trans_indices = np.repeat(trans_indices, repeats=4, axis=2) trans_indices = np.repeat(trans_indices, repeats=3, axis=3) @@ -2493,9 +2438,7 @@ def apply_elf_conformer_selection( self, percentage: float = 2.0, limit: int = 10, - toolkit_registry: Optional[ - Union[ToolkitRegistry, ToolkitWrapper] - ] = GLOBAL_TOOLKIT_REGISTRY, + toolkit_registry: Optional[Union[ToolkitRegistry, ToolkitWrapper]] = GLOBAL_TOOLKIT_REGISTRY, **kwargs, ): """Select a set of diverse conformers from the molecule's conformers with ELF. @@ -2855,7 +2798,7 @@ def to_networkx(self) -> "nx.Graph": """ import networkx as nx - G: nx.classes.graph.Graph = nx.Graph() + G: nx.classes.graph.Graph[int] = nx.Graph() for atom in self.atoms: G.add_node( atom.molecule_atom_index, @@ -2960,9 +2903,9 @@ def _add_atom( atomic_number: int, formal_charge: int, is_aromatic: bool, - stereochemistry: Optional[str] = None, - name: Optional[str] = None, - metadata=None, + stereochemistry: Literal["R", "S", None] = None, + name: str | None = None, + metadata: dict[str, int | str] | None = None, invalidate_cache: bool = True, ) -> int: """ @@ -3030,12 +2973,12 @@ def _add_atom( def _add_bond( self, - atom1, - atom2, - bond_order, - is_aromatic, - stereochemistry=None, - fractional_bond_order=None, + atom1: int | Atom, + atom2: int | Atom, + bond_order: int, + is_aromatic: bool, + stereochemistry: Literal["E", "Z", None] = None, + fractional_bond_order: float | None = None, invalidate_cache: bool = True, ): """ @@ -3078,9 +3021,7 @@ def _add_bond( ) # TODO: Check to make sure bond does not already exist if atom1_atom.is_bonded_to(atom2_atom): - raise BondExistsError( - f"Bond already exists between {atom1_atom} and {atom2_atom})" - ) + raise BondExistsError(f"Bond already exists between {atom1_atom} and {atom2_atom})") bond = Bond( atom1_atom, atom2_atom, @@ -3130,8 +3071,7 @@ def _add_conformer(self, coordinates: Quantity): if not isinstance(coordinates, openmm_unit.Quantity): raise IncompatibleUnitError( - "Unsupported type passed to Molecule._add_conformer setter. " - "Found object of type {type(other)}." + "Unsupported type passed to Molecule._add_conformer setter. Found object of type {type(other)}." ) if not coordinates.unit.is_compatible(openmm_unit.meter): @@ -3149,9 +3089,7 @@ def _add_conformer(self, coordinates: Quantity): f"openmm.unit.Quantity and openff.units.unit.Quantity, found type {type(coordinates)}." ) - tmp_conf = Quantity( - np.zeros(shape=(self.n_atoms, 3), dtype=float), unit.angstrom - ) + tmp_conf = Quantity(np.zeros(shape=(self.n_atoms, 3), dtype=float), unit.angstrom) try: tmp_conf[:] = coordinates # type: ignore[index] except AttributeError as e: @@ -3219,8 +3157,7 @@ def partial_charges(self, charges): if not isinstance(charges, openmm_unit.Quantity): raise IncompatibleUnitError( - "Unsupported type passed to partial_charges setter. " - f"Found object of type {type(charges)}." + f"Unsupported type passed to partial_charges setter. Found object of type {type(charges)}." ) else: @@ -3275,7 +3212,7 @@ def n_impropers(self) -> int: return len(self._impropers) @property - def atoms(self): + def atoms(self) -> list[Atom]: """ Iterate over all Atom objects in the molecule. """ @@ -3373,9 +3310,7 @@ def torsions(self) -> set[tuple[Atom, Atom, Atom, Atom]]: torsions """ self._construct_torsions() - assert ( - self._torsions is not None - ), "_construct_torsions always sets _torsions to a set" + assert self._torsions is not None, "_construct_torsions always sets _torsions to a set" return self._torsions @property @@ -3388,9 +3323,7 @@ def propers(self) -> set[tuple[Atom, Atom, Atom, Atom]]: * Do we need to return a ``Torsion`` object that collects information about fractional bond orders? """ self._construct_torsions() - assert ( - self._propers is not None - ), "_construct_torsions always sets _propers to a set" + assert self._propers is not None, "_construct_torsions always sets _propers to a set" return self._propers @property @@ -3450,11 +3383,7 @@ def smirnoff_impropers(self) -> set[tuple[Atom, Atom, Atom, Atom]]: impropers, amber_impropers """ - return { - improper - for improper in self.impropers - if len(self._bonded_atoms[improper[1]]) == 3 - } + return {improper for improper in self.impropers if len(self._bonded_atoms[improper[1]]) == 3} @property def amber_impropers(self) -> set[tuple[Atom, Atom, Atom, Atom]]: @@ -3485,10 +3414,7 @@ def amber_impropers(self) -> set[tuple[Atom, Atom, Atom, Atom]]: """ self._construct_torsions() - return { - (improper[1], improper[0], improper[2], improper[3]) - for improper in self.smirnoff_impropers - } + return {(improper[1], improper[0], improper[2], improper[3]) for improper in self.smirnoff_impropers} def nth_degree_neighbors(self, n_degrees): """ @@ -3522,9 +3448,7 @@ def nth_degree_neighbors(self, n_degrees): f"path lengths of {n_degrees}." ) else: - return _nth_degree_neighbors_from_graphlike( - graphlike=self, n_degrees=n_degrees - ) + return _nth_degree_neighbors_from_graphlike(graphlike=self, n_degrees=n_degrees) @property def total_charge(self): @@ -3580,7 +3504,7 @@ def to_hill_formula(self) -> str: return self._hill_formula @staticmethod - def _object_to_hill_formula(obj: Union["FrozenMolecule", "nx.Graph"]) -> str: + def _object_to_hill_formula(obj: Union["FrozenMolecule", "nx.Graph[int]"]) -> str: """Take a Molecule or NetworkX graph and generate its Hill formula. This provides a backdoor to the old functionality of Molecule.to_hill_formula, which was a static method that duck-typed inputs of Molecule or graph objects.""" @@ -3592,8 +3516,7 @@ def _object_to_hill_formula(obj: Union["FrozenMolecule", "nx.Graph"]) -> str: return _networkx_graph_to_hill_formula(obj) else: raise TypeError( - "_object_to_hill_formula accepts a NetworkX Graph or OpenFF " - + f"(Frozen)Molecule, not {type(obj)}" + "_object_to_hill_formula accepts a NetworkX Graph or OpenFF " + f"(Frozen)Molecule, not {type(obj)}" ) def chemical_environment_matches( @@ -3812,11 +3735,11 @@ def to_topology(self): @classmethod def from_file( cls: type[FM], - file_path: Union[str, pathlib.Path, TextIO], - file_format=None, - toolkit_registry=GLOBAL_TOOLKIT_REGISTRY, + file_path: str | pathlib.Path | TextIO, + file_format: str | None = None, + toolkit_registry: ToolkitWrapper | ToolkitRegistry = GLOBAL_TOOLKIT_REGISTRY, allow_undefined_stereo: bool = False, - ) -> Union[FM, list[FM]]: + ) -> FM | list[FM]: """ Create one or more molecules from a file @@ -3860,9 +3783,7 @@ def from_file( if isinstance(file_path, pathlib.Path): file_path: str = file_path.as_posix() # type: ignore[no-redef] if not isinstance(file_path, str): - raise ValueError( - "If providing a file-like object for reading molecules, the format must be specified" - ) + raise ValueError("If providing a file-like object for reading molecules, the format must be specified") # Assume that files ending in ".gz" should use their second-to-last suffix for compatibility check # TODO: Will all cheminformatics packages be OK with gzipped files? if file_path[-3:] == ".gz": @@ -3888,9 +3809,7 @@ def from_file( if file_format in query_toolkit.toolkit_file_read_formats: toolkit = query_toolkit break - supported_read_formats[query_toolkit.toolkit_name] = ( - query_toolkit.toolkit_file_read_formats - ) + supported_read_formats[query_toolkit.toolkit_name] = query_toolkit.toolkit_file_read_formats if toolkit is None: msg = ( f"No toolkits in registry can read file {file_path} (format {file_format}). Supported " @@ -4068,12 +3987,7 @@ def from_polymer_pdb( ) coords = Quantity( - np.array( - [ - [*vec3.value_in_unit(openmm_unit.angstrom)] - for vec3 in pdb.getPositions() - ] - ), + np.array([[*vec3.value_in_unit(openmm_unit.angstrom)] for vec3 in pdb.getPositions()]), unit.angstrom, ) offmol.add_conformer(coords) @@ -4121,19 +4035,19 @@ def _to_xyz_file(self, file_path: Union[str, IO[str]]): # If we do not have a conformer make one with all zeros if not self._conformers: - conformers: list[Quantity] = [ - Quantity(np.zeros((self.n_atoms, 3), dtype=float), unit.angstrom) - ] + conformers: list[Quantity] = [Quantity(np.zeros((self.n_atoms, 3), dtype=float), unit.angstrom)] else: conformers = self._conformers if len(conformers) == 1: end: Union[str, int] = "" + def title(frame): return f"{self.name if self.name != '' else self.hill_formula}{frame}\n" else: end = 1 + def title(frame): return f"{self.name if self.name != '' else self.hill_formula} Frame {frame}\n" @@ -4148,9 +4062,7 @@ def title(frame): xyz_data.write(f"{self.n_atoms}\n" + title(end)) for j, atom_coords in enumerate(geometry.m_as(unit.angstrom)): # type: ignore[arg-type] x, y, z = atom_coords - xyz_data.write( - f"{SYMBOLS[self.atoms[j].atomic_number]} {x: .10f} {y: .10f} {z: .10f}\n" - ) + xyz_data.write(f"{SYMBOLS[self.atoms[j].atomic_number]} {x: .10f} {y: .10f} {z: .10f}\n") # now we up the frame count end = i + 1 @@ -4215,9 +4127,7 @@ def to_file(self, file_path, file_format, toolkit_registry=GLOBAL_TOOLKIT_REGIST if toolkit is None: supported_formats = {} for _toolkit in toolkit_registry.registered_toolkits: - supported_formats[_toolkit.toolkit_name] = ( - _toolkit.toolkit_file_write_formats - ) + supported_formats[_toolkit.toolkit_name] = _toolkit.toolkit_file_write_formats raise ValueError( f"The requested file format ({file_format}) is not available from any of the installed toolkits " f"(supported formats: {supported_formats})" @@ -4228,9 +4138,7 @@ def to_file(self, file_path, file_format, toolkit_registry=GLOBAL_TOOLKIT_REGIST else: toolkit.to_file_obj(self, file_path, file_format) - def enumerate_tautomers( - self, max_states=20, toolkit_registry=GLOBAL_TOOLKIT_REGISTRY - ): + def enumerate_tautomers(self, max_states=20, toolkit_registry=GLOBAL_TOOLKIT_REGISTRY): """ Enumerate the possible tautomers of the current molecule @@ -4249,14 +4157,10 @@ def enumerate_tautomers( """ if isinstance(toolkit_registry, ToolkitRegistry): - molecules = toolkit_registry.call( - "enumerate_tautomers", molecule=self, max_states=max_states - ) + molecules = toolkit_registry.call("enumerate_tautomers", molecule=self, max_states=max_states) elif isinstance(toolkit_registry, ToolkitWrapper): - molecules = toolkit_registry.enumerate_tautomers( - self, max_states=max_states - ) + molecules = toolkit_registry.enumerate_tautomers(self, max_states=max_states) else: raise InvalidToolkitRegistryError( @@ -4427,9 +4331,7 @@ def to_rdkit( if isinstance(toolkit_registry, ToolkitWrapper): return toolkit_registry.to_rdkit(self, aromaticity_model=aromaticity_model) # type: ignore[attr-defined] else: - return toolkit_registry.call( - "to_rdkit", self, aromaticity_model=aromaticity_model - ) + return toolkit_registry.call("to_rdkit", self, aromaticity_model=aromaticity_model) @classmethod @OpenEyeToolkitWrapper.requires_toolkit() @@ -4469,9 +4371,7 @@ def from_openeye( """ toolkit = OpenEyeToolkitWrapper() - molecule = toolkit.from_openeye( - oemol, allow_undefined_stereo=allow_undefined_stereo, _cls=cls - ) + molecule = toolkit.from_openeye(oemol, allow_undefined_stereo=allow_undefined_stereo, _cls=cls) return molecule @requires_package("qcelemental") @@ -4535,25 +4435,13 @@ def to_qcschema(self, multiplicity=1, conformer=0, extras=None): # Gather the required qcschema data charge = self.total_charge.m_as(unit.elementary_charge) - connectivity = [ - (bond.atom1_index, bond.atom2_index, bond.bond_order) for bond in self.bonds - ] + connectivity = [(bond.atom1_index, bond.atom2_index, bond.bond_order) for bond in self.bonds] symbols = [SYMBOLS[atom.atomic_number] for atom in self.atoms] if extras is not None: - extras["canonical_isomeric_explicit_hydrogen_mapped_smiles"] = ( - self.to_smiles(mapped=True) - ) + extras["canonical_isomeric_explicit_hydrogen_mapped_smiles"] = self.to_smiles(mapped=True) else: - extras = { - "canonical_isomeric_explicit_hydrogen_mapped_smiles": self.to_smiles( - mapped=True - ) - } - identifiers = { - "canonical_isomeric_explicit_hydrogen_mapped_smiles": self.to_smiles( - mapped=True - ) - } + extras = {"canonical_isomeric_explicit_hydrogen_mapped_smiles": self.to_smiles(mapped=True)} + identifiers = {"canonical_isomeric_explicit_hydrogen_mapped_smiles": self.to_smiles(mapped=True)} schema_dict = { "symbols": symbols, @@ -4660,9 +4548,7 @@ def from_mapped_smiles( ) if len(mapping) != offmol.n_atoms: - raise SmilesParsingError( - "The mapped smiles does not contain enough indexes to remap the molecule." - ) + raise SmilesParsingError("The mapped smiles does not contain enough indexes to remap the molecule.") # remap the molecule using the atom map found in the smiles # the order is mapping = dict[current_index: new_index] @@ -4788,28 +4674,18 @@ def from_qcschema( # so we don't need to cast this to list mol_dicts = qca_object.get("initial_molecules") if not mol_dicts: - raise InvalidQCInputError( - f"Unable to find molecule information in qcschema input. {qca_object=}" - ) + raise InvalidQCInputError(f"Unable to find molecule information in qcschema input. {qca_object=}") first_cmiles = None for mol_dict in mol_dicts: # Entries sometimes have their cmiles here - cmiles = qca_object.get("attributes", {}).get( - "canonical_isomeric_explicit_hydrogen_mapped_smiles" - ) + cmiles = qca_object.get("attributes", {}).get("canonical_isomeric_explicit_hydrogen_mapped_smiles") if not cmiles: - cmiles = mol_dict.get("identifiers", {}).get( - "canonical_isomeric_explicit_hydrogen_mapped_smiles" - ) + cmiles = mol_dict.get("identifiers", {}).get("canonical_isomeric_explicit_hydrogen_mapped_smiles") if not cmiles: - cmiles = mol_dict.get("extras", {}).get( - "canonical_isomeric_explicit_hydrogen_mapped_smiles" - ) + cmiles = mol_dict.get("extras", {}).get("canonical_isomeric_explicit_hydrogen_mapped_smiles") if not cmiles: - raise MissingCMILESError( - f"Unable to find CMILES in qcschema input molecule. {mol_dict=}" - ) + raise MissingCMILESError(f"Unable to find CMILES in qcschema input molecule. {mol_dict=}") if first_cmiles is None: first_cmiles = cmiles offmol = cls.from_mapped_smiles( @@ -4824,9 +4700,7 @@ def from_qcschema( f"{first_cmiles} != {cmiles} when iterating over molecules for " f"input {qca_object}" ) - geometry = Quantity( - np.array(mol_dict["geometry"], float).reshape(-1, 3), unit.bohr - ) + geometry = Quantity(np.array(mol_dict["geometry"], float).reshape(-1, 3), unit.bohr) offmol._add_conformer(geometry.to(unit.angstrom)) # If there's a QCA ID for this QC molecule, store it in the OFF molecule with reference to # its corresponding conformer @@ -4890,9 +4764,7 @@ def from_pdb_and_smiles( ) toolkit = RDKitToolkitWrapper() - return toolkit.from_pdb_and_smiles( - file_path, smiles, allow_undefined_stereo, _cls=cls, name=name - ) + return toolkit.from_pdb_and_smiles(file_path, smiles, allow_undefined_stereo, _cls=cls, name=name) def canonical_order_atoms(self, toolkit_registry=GLOBAL_TOOLKIT_REGISTRY): """ @@ -4990,9 +4862,7 @@ def remap( """ # make sure the size of the mapping matches the current molecule - if len(mapping_dict) > self.n_atoms or ( - len(mapping_dict) < self.n_atoms and not partial - ): + if len(mapping_dict) > self.n_atoms or (len(mapping_dict) < self.n_atoms and not partial): raise RemapIndexError( f"The number of mapping indices ({len(mapping_dict)}) does not " + f"match the number of atoms in this molecule ({self.n_atoms})" @@ -5009,15 +4879,9 @@ def remap( # Make sure that there were no duplicate indices if len(new_to_cur) != len(cur_to_new): - raise RemapIndexError( - "There must be no duplicate source or destination indices in" - + " mapping_dict" - ) + raise RemapIndexError("There must be no duplicate source or destination indices in" + " mapping_dict") - if any( - not (isinstance(i, int) and 0 <= i < self.n_atoms) - for i in [*new_to_cur, *cur_to_new] - ): + if any(not (isinstance(i, int) and 0 <= i < self.n_atoms) for i in [*new_to_cur, *cur_to_new]): raise RemapIndexError( f"All indices in a mapping_dict for a molecule with {self.n_atoms}" + f" atoms must be integers between 0 and {self.n_atoms - 1}" @@ -5048,13 +4912,11 @@ def remap( for i in range(self.n_atoms): # get the old atom info old_atom = self._atoms[new_to_cur[i]] - new_molecule._add_atom(**old_atom.to_dict()) + new_molecule._add_atom(**old_atom.to_dict()) # type:ignore # this is the first time we access the mapping; catch an index error # here corresponding to mapping that starts from 0 or higher except (KeyError, IndexError): - raise RemapIndexError( - f"The mapping supplied is missing a destination index for atom {i}" - ) + raise RemapIndexError(f"The mapping supplied is missing a destination index for atom {i}") # add the bonds but with atom indexes in a sorted ascending order for bond in self._bonds: @@ -5062,21 +4924,17 @@ def remap( bond_dict = bond.to_dict() bond_dict["atom1"] = atoms[0] bond_dict["atom2"] = atoms[1] - new_molecule._add_bond(**bond_dict) + new_molecule._add_bond(**bond_dict) # type:ignore # we can now resort the bonds - sorted_bonds = sorted( - new_molecule.bonds, key=operator.attrgetter("atom1_index", "atom2_index") - ) + sorted_bonds = sorted(new_molecule.bonds, key=operator.attrgetter("atom1_index", "atom2_index")) new_molecule._bonds = sorted_bonds # remap the charges if self.partial_charges is not None: new_charges = np.zeros(self.n_atoms) for i in range(self.n_atoms): - new_charges[i] = self.partial_charges[new_to_cur[i]].m_as( - unit.elementary_charge - ) + new_charges[i] = self.partial_charges[new_to_cur[i]].m_as(unit.elementary_charge) new_molecule.partial_charges = new_charges * unit.elementary_charge # remap the conformers, there can be more than one @@ -5091,12 +4949,9 @@ def remap( new_molecule._properties = deepcopy(self._properties) # remap the atom map - if "atom_map" in new_molecule.properties and isinstance( - new_molecule.properties["atom_map"], dict - ): + if "atom_map" in new_molecule.properties and isinstance(new_molecule.properties["atom_map"], dict): new_molecule.properties["atom_map"] = { - cur_to_new.get(k, k): v - for k, v in new_molecule.properties["atom_map"].items() + cur_to_new.get(k, k): v for k, v in new_molecule.properties["atom_map"].items() } return new_molecule @@ -5141,9 +4996,7 @@ def to_openeye( self, aromaticity_model=aromaticity_model ) else: - return toolkit_registry.call( - "to_openeye", self, aromaticity_model=aromaticity_model - ) + return toolkit_registry.call("to_openeye", self, aromaticity_model=aromaticity_model) def _construct_angles(self) -> None: """ @@ -5262,10 +5115,7 @@ def get_bond_between(self, i: Union[int, Atom], j: Union[int, Atom]) -> Bond: atom_i = i atom_j = j else: - raise TypeError( - "Invalid input passed to get_bond_between(). Expected ints or Atoms, " - f"got {j} and {j}." - ) + raise TypeError(f"Invalid input passed to get_bond_between(). Expected ints or Atoms, got {j} and {j}.") for bond in atom_i.bonds: for atom in bond.atoms: @@ -5347,7 +5197,7 @@ def add_atom( atomic_number: int, formal_charge: int, is_aromatic: bool, - stereochemistry: Optional[str] = None, + stereochemistry: Literal["R", "S", None] = None, name: Optional[str] = None, metadata: Optional[dict[str, Union[int, str]]] = None, ) -> int: @@ -5413,7 +5263,7 @@ def add_bond( atom2: Union[int, "Atom"], bond_order: int, is_aromatic: bool, - stereochemistry: Optional[str] = None, + stereochemistry: Literal["E", "Z", None] = None, fractional_bond_order: Optional[float] = None, ) -> int: """ @@ -5546,9 +5396,7 @@ def visualize( raise MissingOptionalDependencyError("nglview") signature = inspect.signature(Molecule.visualize).parameters - if (width != signature["width"].default) or ( - height != signature["height"].default - ): + if (width != signature["width"].default) or (height != signature["height"].default): warnings.warn( f"Arguments `width` and `height` are ignored with {backend=}." f"Found non-default values {width=} and {height=}", @@ -5557,8 +5405,7 @@ def visualize( if self.conformers is None: raise MissingConformersError( - "Visualizing with NGLview requires that the molecule has " - f"conformers, found {self.conformers=}" + f"Visualizing with NGLview requires that the molecule has conformers, found {self.conformers=}" ) else: @@ -5626,9 +5473,7 @@ def visualize( oemol = self.to_openeye() - opts = oedepict.OE2DMolDisplayOptions( - width, height, oedepict.OEScale_AutoScale - ) + opts = oedepict.OE2DMolDisplayOptions(width, height, oedepict.OEScale_AutoScale) if show_all_hydrogens: opts.SetHydrogenStyle(oedepict.OEHydrogenStyle_ImplicitAll) @@ -5668,9 +5513,7 @@ def perceive_residues( """ # Read substructure dictionary file if not substructure_file_path: - substructure_file_path = get_data_file_path( - "proteins/aa_residues_substructures_with_caps.json" - ) + substructure_file_path = get_data_file_path("proteins/aa_residues_substructures_with_caps.json") with open(substructure_file_path) as subfile: substructure_dictionary = json.load(subfile) @@ -5683,10 +5526,8 @@ def perceive_residues( for res_name, inner_dict in substructure_dictionary.items(): for smarts in inner_dict.keys(): smarts_no_chirality = smarts.replace("@", "") # remove @ in smarts - substructure_dictionary_no_chirality[res_name][ - smarts_no_chirality - ] = substructure_dictionary_no_chirality[res_name].pop( - smarts + substructure_dictionary_no_chirality[res_name][smarts_no_chirality] = ( + substructure_dictionary_no_chirality[res_name].pop(smarts) ) # update key # replace with the new substructure dictionary substructure_dictionary = substructure_dictionary_no_chirality @@ -5714,13 +5555,9 @@ def perceive_residues( this_match_set = all_matches[match_idx]["atom_idxs_set"] this_match_set_size = len(this_match_set) for match_before_this_idx in range(match_idx): - match_before_this_set = all_matches[match_before_this_idx][ - "atom_idxs_set" - ] + match_before_this_set = all_matches[match_before_this_idx]["atom_idxs_set"] match_before_this_set_size = len(match_before_this_set) - n_overlapping_atoms = len( - this_match_set.intersection(match_before_this_set) - ) + n_overlapping_atoms = len(this_match_set.intersection(match_before_this_set)) if n_overlapping_atoms > 0: if match_before_this_set_size < this_match_set_size: match_idxs_to_delete.add(match_before_this_idx) @@ -5736,14 +5573,10 @@ def perceive_residues( # Now the matches have been deduplicated and de-subsetted for residue_num, match_dict in enumerate(all_matches): for smarts_idx, atom_idx in enumerate(match_dict["atom_idxs"]): - self.atoms[atom_idx].metadata["residue_name"] = match_dict[ - "residue_name" - ] + self.atoms[atom_idx].metadata["residue_name"] = match_dict["residue_name"] self.atoms[atom_idx].metadata["residue_number"] = str(residue_num + 1) self.atoms[atom_idx].metadata["insertion_code"] = " " - self.atoms[atom_idx].metadata["atom_name"] = match_dict["atom_names"][ - smarts_idx - ] + self.atoms[atom_idx].metadata["atom_name"] = match_dict["atom_names"][smarts_idx] # Now add the residue hierarchy scheme self._add_residue_hierarchy_scheme() @@ -5767,7 +5600,7 @@ def _ipython_display_(self): # pragma: no cover pass -def _networkx_graph_to_hill_formula(graph: "nx.Graph") -> str: +def _networkx_graph_to_hill_formula(graph: "nx.Graph[int]") -> str: """ Convert a NetworkX graph to a Hill formula. @@ -5788,7 +5621,7 @@ def _networkx_graph_to_hill_formula(graph: "nx.Graph") -> str: raise ValueError("The graph must be a NetworkX graph.") atom_nums = list(dict(graph.nodes(data="atomic_number", default=1)).values()) - return _atom_nums_to_hill_formula(atom_nums) + return _atom_nums_to_hill_formula(atom_nums) # type:ignore[arg-type] def _atom_nums_to_hill_formula(atom_nums: list[int]) -> str: @@ -5824,9 +5657,7 @@ def _atom_nums_to_hill_formula(atom_nums: list[int]) -> str: def _nth_degree_neighbors_from_graphlike( graphlike: MoleculeLike, n_degrees: int, -) -> Generator[ - Union[tuple[Atom, Atom], tuple["_SimpleAtom", "_SimpleAtom"]], None, None -]: +) -> Generator[Union[tuple[Atom, Atom], tuple["_SimpleAtom", "_SimpleAtom"]], None, None]: """ Given a graph-like object, return a tuple of the nth degree neighbors of each atom. @@ -5915,9 +5746,7 @@ def __init__( The name of the iterator that will be exposed to access the hierarchy elements generated by this scheme """ - if (type(uniqueness_criteria) is not list) and ( - type(uniqueness_criteria) is not tuple - ): + if (type(uniqueness_criteria) is not list) and (type(uniqueness_criteria) is not tuple): raise TypeError( f"'uniqueness_criteria' kwarg must be a list or a tuple of strings," f" received {uniqueness_criteria!r} " @@ -5951,9 +5780,7 @@ def to_dict(self) -> dict: return_dict: dict[str, Union[str, Sequence[Union[str, int, dict]]]] = dict() return_dict["uniqueness_criteria"] = self.uniqueness_criteria return_dict["iterator_name"] = self.iterator_name - return_dict["hierarchy_elements"] = [ - e.to_dict() for e in self.hierarchy_elements - ] + return_dict["hierarchy_elements"] = [e.to_dict() for e in self.hierarchy_elements] return return_dict def perceive_hierarchy(self): @@ -5975,9 +5802,7 @@ def perceive_hierarchy(self): self.hierarchy_elements = list() # Determine which atoms should get added to which HierarchyElements - hier_eles_to_add: defaultdict[tuple[Union[int, str]], list[Atom]] = ( - defaultdict(list) - ) + hier_eles_to_add: defaultdict[tuple[Union[int, str]], list[Atom]] = defaultdict(list) for atom in self.parent.atoms: _atom_key = list() for field_key in self.uniqueness_criteria: @@ -6084,7 +5909,7 @@ class HierarchyElement: def __init__( self, scheme: HierarchyScheme, - identifier: tuple[Union[str, int]], + identifier: tuple[str | int, ...], atom_indices: Sequence[int], ): """ @@ -6104,12 +5929,10 @@ def __init__( self.scheme = scheme self.identifier = identifier self.atom_indices = deepcopy(atom_indices) - for id_component, uniqueness_component in zip( - identifier, scheme.uniqueness_criteria - ): + for id_component, uniqueness_component in zip(identifier, scheme.uniqueness_criteria): setattr(self, uniqueness_component, id_component) - def to_dict(self) -> dict[str, Union[tuple[Union[str, int]], Sequence[int]]]: + def to_dict(self) -> dict[str, tuple[str | int, ...] | Sequence[int]]: """Serialize this object to a basic dict of strings and lists of ints. Keys and values align with parameters used to initialize the :class:`HierarchyElement` class. @@ -6180,9 +6003,7 @@ def generate_unique_atom_names(self, suffix: str = "x"): return _generate_unique_atom_names(self, suffix) -def _has_unique_atom_names( - obj: Union[FrozenMolecule, "_SimpleMolecule", HierarchyElement] -) -> bool: +def _has_unique_atom_names(obj: Union[FrozenMolecule, "_SimpleMolecule", HierarchyElement]) -> bool: """``True`` if the object has unique atom names, ``False`` otherwise.""" unique_atom_names = set([atom.name for atom in obj.atoms]) if len(unique_atom_names) < obj.n_atoms: @@ -6190,9 +6011,7 @@ def _has_unique_atom_names( return True -def _generate_unique_atom_names( - obj: Union[FrozenMolecule, HierarchyElement], suffix: str = "x" -): +def _generate_unique_atom_names(obj: Union[FrozenMolecule, HierarchyElement], suffix: str = "x"): """ Generate unique atom names from the element symbol and count. diff --git a/openff/toolkit/topology/topology.py b/openff/toolkit/topology/topology.py index be98c001d..245c3f70e 100644 --- a/openff/toolkit/topology/topology.py +++ b/openff/toolkit/topology/topology.py @@ -525,7 +525,7 @@ def n_unique_molecules(self) -> int: @classmethod def from_molecules( cls, - molecules: Union[MoleculeLike, list[MoleculeLike]], + molecules: MoleculeLike | Iterable[MoleculeLike], ) -> "Topology": """ Create a new Topology object containing one copy of each of the specified molecule(s). @@ -703,7 +703,7 @@ def n_molecules(self) -> int: return len(self._molecules) @property - def molecules(self) -> Generator[MoleculeLike, None, None]: + def molecules(self) -> Iterator[MoleculeLike]: """Returns an iterator over all the Molecules in this Topology Returns @@ -1085,7 +1085,7 @@ def chemical_environment_matches( topology_atom_indices = [] for molecule_atom_index in match: atom = mol_instance.atom(atom_map[molecule_atom_index]) - topology_atom_indices.append(self.atom_index(atom)) + topology_atom_indices.append(self.atom_index(atom)) # type: ignore[arg-type] environment_match = Topology._ChemicalEnvironmentMatch( tuple(match), unique_mol, tuple(topology_atom_indices) @@ -1810,10 +1810,10 @@ def from_pdb( ) for off_atom, atom in zip([*topology.atoms], pdb.topology.atoms()): - off_atom.metadata["residue_name"] = atom.residue.name - off_atom.metadata["residue_number"] = atom.residue.id - off_atom.metadata["insertion_code"] = atom.residue.insertionCode - off_atom.metadata["chain_id"] = atom.residue.chain.id + off_atom.metadata["residue_name"] = atom.residue.name # type:ignore[attr-defined] + off_atom.metadata["residue_number"] = atom.residue.id # type:ignore[attr-defined] + off_atom.metadata["insertion_code"] = atom.residue.insertionCode # type:ignore[attr-defined] + off_atom.metadata["chain_id"] = atom.residue.chain.id # type:ignore[attr-defined] off_atom.name = atom.name for offmol in topology.molecules: @@ -2135,7 +2135,7 @@ def clear_positions(self): for molecule in self.molecules: molecule._conformers = None - def set_positions(self, array: Quantity): + def set_positions(self, array: Quantity) -> None: """ Set the positions in a topology by copying from a single (n, 3) array. @@ -2343,7 +2343,8 @@ def atom(self, atom_topology_index: int) -> "Atom": if next_molecule_start_index > atom_topology_index: atom_molecule_index = atom_topology_index - this_molecule_start_index # NOTE: the index here should still be in the topology index order, NOT the reference molecule's - return molecule.atom(atom_molecule_index) + # can Molecule.atom be a _SimpleAtom? + return molecule.atom(atom_molecule_index) # type: ignore[return-value] this_molecule_start_index += molecule.n_atoms raise AtomNotInTopologyError( diff --git a/openff/toolkit/utils/builtin_wrapper.py b/openff/toolkit/utils/builtin_wrapper.py index 5cc6daa34..2666e163a 100644 --- a/openff/toolkit/utils/builtin_wrapper.py +++ b/openff/toolkit/utils/builtin_wrapper.py @@ -131,7 +131,7 @@ def assign_partial_charges( partial_charges = [0.0] * molecule.n_atoms elif partial_charge_method == "formal_charge": - partial_charges = [float(atom.formal_charge.m) for atom in molecule.atoms] + partial_charges = [float(atom.formal_charge.m) for atom in molecule.atoms] # type: ignore molecule.partial_charges = Quantity(partial_charges, unit.elementary_charge) diff --git a/openff/toolkit/utils/rdkit_wrapper.py b/openff/toolkit/utils/rdkit_wrapper.py index ebf7eb8be..21d783cca 100644 --- a/openff/toolkit/utils/rdkit_wrapper.py +++ b/openff/toolkit/utils/rdkit_wrapper.py @@ -3259,7 +3259,7 @@ def _detect_undefined_stereo( atom1, atom2 = bond.GetBeginAtom(), bond.GetEndAtom() msg += ( f" - Bond {undefined_bond_idx} (atoms {atom1.GetIdx()}-{atom2.GetIdx()} of element " - "({atom1.GetSymbol()}-{atom2.GetSymbol()})\n" + f"({atom1.GetSymbol()}-{atom2.GetSymbol()})\n" ) raise UndefinedStereochemistryError(err_msg_prefix + msg)