Source code for stk.molecular.writers.mdl_mol

"""
Mol Writer
==========

"""

from __future__ import annotations

import typing

from ...utilities import OneOrMany
from ..molecules import Molecule


[docs]class MolWriter: """ A writer class for V3000 ``.mol`` files. Examples: *Writing to a File* .. testcode:: writing-to-a-file import stk bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()]) writer = stk.MolWriter() writer.write(molecule=bb1, path='bb1.mol') .. testcode:: writing-to-a-file :hide: import os assert os.path.exists('bb1.mol') .. testcleanup:: writing-to-a-file os.remove('bb1.mol') """ def _write_content( self, molecule: Molecule, atom_ids: typing.Optional[OneOrMany[int]], ) -> str: if atom_ids is None: atom_ids = range(molecule.get_num_atoms()) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) atom_lines = [] coords = molecule.get_position_matrix() # This set gets used by bonds. id_map = {} for new_atom_id, old_atom_id in enumerate(atom_ids, 1): id_map[old_atom_id] = new_atom_id x, y, z = (i for i in coords[old_atom_id]) (atom,) = molecule.get_atoms(atom_ids=old_atom_id) symbol = atom.__class__.__name__ charge_ = atom.get_charge() charge = f" CHG={charge_}" if charge_ else "" atom_lines.append( f"M V30 {new_atom_id} {symbol} {x:.4f} " f"{y:.4f} {z:.4f} 0{charge}\n" ) atom_block = "".join(atom_lines) bond_lines: list[str] = [] for bond in molecule.get_bonds(): a1 = bond.get_atom1().get_id() a2 = bond.get_atom2().get_id() if a1 in id_map and a2 in id_map: bond_lines.append( f"M V30 {len(bond_lines)+1} " f"{int(bond.get_order())} " f"{id_map[a1]} {id_map[a2]}\n" ) num_bonds = len(bond_lines) bond_block = "".join(bond_lines) return ( "\n" " RDKit 3D\n" "\n" " 0 0 0 0 0 0 0 0 0 0999 V3000\n" "M V30 BEGIN CTAB\n" f"M V30 COUNTS {len(id_map)} {num_bonds} 0 0 0\n" "M V30 BEGIN ATOM\n" f"{atom_block}" "M V30 END ATOM\n" "M V30 BEGIN BOND\n" f"{bond_block}" "M V30 END BOND\n" "M V30 END CTAB\n" "M END\n" "\n" "$$$$\n" )
[docs] def to_string( self, molecule: Molecule, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> str: """ Get a V3000 ``.mol`` file format string of `molecule`. Parameters: molecule: Molecule to write to V3000 ``.mol`` format. atom_ids: The atom ids of atoms to write. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. Returns: String in V3000 ``.mol`` file format. """ return self._write_content(molecule, atom_ids)
[docs] def write( self, molecule: Molecule, path: str, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> None: """ Write `molecule` to V3000 ``.mol`` file format. Parameters: molecule: Molecule to write to V3000 ``.mol`` format. path: The full path to the file being written. atom_ids: The atom ids of atoms to write. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. """ with open(path, "w") as f: f.write(self._write_content(molecule, atom_ids))