"""Import and export: SMILES, MOL, SDF, RXN, CDXML, InChI, names, KET and XYZ.
Import turns a file into molecules (and a reaction, when the format carries
one), writes them onto the canvas as KET and derives the document from it,
so an imported structure is drawn exactly like a drawn one. An XYZ file, as
vibe-view exports one, keeps its coordinates as each species' 3D structure. Export reads the
document: structures keep the coordinates they have on the canvas.
Formats a document cannot be expressed in are refused with the reason,
never approximated: an RXN file needs a reaction, PNG of a reaction scheme
needs a renderer this version does not have. CDXML is read and written as a
whole drawing by :mod:`vibesketch.document.cdxml`; binary CDX is read through
it (:mod:`vibesketch.document.cdx`).
"""
from __future__ import annotations
import contextlib
import io
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal
from vibesketch.document.derive import Derivation, derive
from vibesketch.document.io import dumps, new_document
from vibesketch.document.ket import read_scene, stereo_reason, write_scene
from vibesketch.document.model import Canvas, Document
__all__ = [
"EXPORT_FORMATS",
"IMPORT_FORMATS",
"TEXT_FORMATS",
"ExportRefused",
"ImportFailed",
"Imported",
"Payload",
"export",
"format_for_path",
"import_into",
"read_structures",
"species_mols",
]
ImportFormat = Literal[
"smiles",
"mol",
"sdf",
"rxn",
"cdxml",
"cdx",
"inchi",
"name",
"ket",
"vsk",
"qvf",
"xyz",
"orca",
"orca-hess",
"peptide",
"dna",
"rna",
"helm",
]
ExportFormat = Literal[
"qvf",
"vsk",
"ket",
"mol",
"sdf",
"rxn",
"smiles",
"inchi",
"svg",
"png",
"pdf",
"chemfig",
"helm",
"fasta",
]
IMPORT_FORMATS: tuple[str, ...] = (
"smiles", "mol", "sdf", "rxn", "cdxml", "cdx", "inchi", "name", "ket", "vsk", "qvf", "xyz",
"orca", "orca-hess", "peptide", "dna", "rna", "helm",
) # fmt: skip
#: Formats whose data is short text, which the CLI and the browser take as typed.
TEXT_FORMATS: tuple[str, ...] = ("smiles", "inchi", "name", "peptide", "dna", "rna", "helm")
#: Calculation files: atoms and coordinates, whose bonds are inferred (#86).
CALCULATIONS: tuple[str, ...] = ("orca", "orca-hess")
#: QVF first: the format of every file in the vibe family.
EXPORT_FORMATS: tuple[str, ...] = (
"qvf", "vsk", "ket", "mol", "sdf", "rxn", "smiles", "inchi", "svg", "png", "pdf",
"eps", "emf", "cdxml", "chemfig", "helm", "fasta",
) # fmt: skip
#: Formats that carry a whole document, not just structures.
WHOLE_DOCUMENTS: tuple[str, ...] = ("qvf", "vsk")
_SUFFIXES = {
".smi": "smiles", ".smiles": "smiles", ".mol": "mol", ".sdf": "sdf", ".sd": "sdf",
".rxn": "rxn", ".cdxml": "cdxml", ".inchi": "inchi", ".ket": "ket", ".vsk": "vsk",
".qvf": "qvf", ".xyz": "xyz", ".out": "orca", ".hess": "orca-hess", ".cdx": "cdx",
".helm": "helm",
} # fmt: skip
MEDIA_TYPES = {
"qvf": "application/vnd.vibe-qc.qvf+zip", # registered in the qvf repository, v0.2
"vsk": "application/vnd.vibesketch+json",
"ket": "application/json",
"mol": "chemical/x-mdl-molfile",
"sdf": "chemical/x-mdl-sdfile",
"rxn": "chemical/x-mdl-rxnfile",
"cdxml": "chemical/x-cdxml",
"cdx": "chemical/x-cdx",
"smiles": "chemical/x-daylight-smiles",
"inchi": "chemical/x-inchi",
"svg": "image/svg+xml",
"png": "image/png",
"pdf": "application/pdf",
"eps": "application/postscript",
"emf": "image/emf", # RFC 7903
"chemfig": "application/x-tex",
# Neither has a registered media type.
"helm": "text/plain",
"fasta": "text/plain",
}
EXTENSIONS = {"smiles": "smi", "inchi": "inchi", "chemfig": "tex"}
[docs]
class ImportFailed(ValueError):
"""The data could not be read in the stated format."""
[docs]
class ExportRefused(ValueError):
"""The document cannot be expressed in the requested format."""
[docs]
@dataclass
class Imported:
molecules: list[Any] = field(default_factory=list)
reaction: tuple[list[Any], list[Any], list[Any]] | None = None
conditions: list[str] = field(default_factory=list)
#: A whole canvas (KET or a .vsk's), used as it is.
ket: dict[str, Any] | None = None
#: What did not carry over from the file.
warnings: list[str] = field(default_factory=list)
#: Whether the molecules were built from sequences, which their species record (#98).
sequenced: bool = False
#: R-group definitions read from an R-group file, drawn beside the molecules (#90).
rgroups: list[Any] = field(default_factory=list)
[docs]
@dataclass(frozen=True)
class Payload:
content: bytes
media_type: str
extension: str
def _text(data: str | bytes) -> str:
return data.decode("utf-8") if isinstance(data, bytes) else data
def _sanitized(mol: Any, what: str) -> Any:
from rdkit import Chem
if mol is None:
raise ImportFailed(f"RDKit could not read the {what}")
try:
Chem.SanitizeMol(mol)
except Exception as error:
raise ImportFailed(f"the {what} is not a valid structure: {error}") from error
from vibesketch.chem.species import inert_pair_radicals
inert_pair_radicals(mol) # Pb²⁺ gains no radicals the file did not hold (#147)
# Honor explicit CTfile relative declarations before converting to KET,
# whose normalized species model represents individual stereoisomers.
Chem.TranslateChiralFlagToStereoGroups(mol)
if any(
group.GetGroupType() != Chem.StereoGroupType.STEREO_ABSOLUTE
for group in mol.GetStereoGroups()
):
raise ImportFailed(
f"the {what} uses relative or mixture stereochemistry; "
"choose an explicit absolute stereoisomer before importing for calculation"
)
return mol
def _reaction(rxn: Any, what: str) -> tuple[list[Any], list[Any], list[Any]]:
def mols(templates: Any) -> list[Any]:
return [_sanitized(m, what) for m in templates]
return (
mols(rxn.GetReactants()),
mols(rxn.GetAgents()),
mols(rxn.GetProducts()),
)
def _check_cx_stereo(text: str) -> None:
"""Refuse the legacy relative flag, which RDKit accepts but ignores."""
parts = text.split(maxsplit=1)
if len(parts) < 2 or not parts[1].startswith("|"):
return
extension = parts[1].split("|", 2)[1]
# Coordinates and atom-label text are not top-level CXSMILES fields.
extension = re.sub(r"\$[^$]*\$|\([^)]*\)", "", extension)
# Chemaxon CXSMILES: r or r:<fragment indices> declares relative stereo.
# https://docs.chemaxon.com/latest/formats_chemaxon-extended-smiles-and-smarts-cxsmiles-and-cxsmarts.html
# RDKit also accepts whitespace before the field delimiter or closing pipe.
if re.search(r"(?:^|,)\s*r\s*(?=[:,]|$)", extension):
raise ImportFailed(
"CXSMILES declares relative or mixture stereochemistry with the r flag; "
"choose an explicit absolute stereoisomer before importing for calculation"
)
#: ``key=value`` pairs an XYZ comment line may carry, as vibe-sketch's XYZ writer puts them.
_XYZ_STATED = re.compile(r"\b(charge|multiplicity|mult)\s*=\s*([+-]?\d+)\b", re.IGNORECASE)
def _element(token: str, where: str) -> str:
"""An element symbol from an XYZ atom label: ``C``, ``cl``, ``C12`` or an atomic number."""
from vibesketch.elements import SYMBOLS
if token.isdigit():
number = int(token)
if 1 <= number <= len(SYMBOLS):
return SYMBOLS[number - 1]
else:
letters = token.rstrip("0123456789")
symbol = letters[:1].upper() + letters[1:].lower()
if symbol in SYMBOLS:
return symbol
raise ImportFailed(f"{where}: {token!r} is no element")
def _xyz(data: str | bytes) -> Derivation:
"""The frames of an XYZ file as species with their 3D structures.
A frame is a line with the atom count, a comment line, and a line per
atom with its symbol and x, y, z in angstrom. ``charge=`` and
``multiplicity=`` on the comment line are taken; without them a frame is
neutral. Frames of the same molecule, as a trajectory has, give one
species with the last frame's coordinates.
"""
from vibesketch.document.coordinates import NothingDrawn, Structure, from_structures
lines = _text(data).splitlines()
frames: list[Structure] = []
row = 0
while row < len(lines):
if not lines[row].strip():
row += 1
continue
where = f"frame {len(frames) + 1}"
head = lines[row].split()
if not head[0].isdigit():
raise ImportFailed(f"line {row + 1}: an XYZ frame starts with its atom count")
count = int(head[0])
atoms = lines[row + 2 : row + 2 + count]
if count == 0 or len(atoms) < count:
raise ImportFailed(f"{where}: {count} atoms announced, {len(atoms)} given")
symbols: list[str] = []
positions: list[list[float]] = []
for number, line in enumerate(atoms, start=1):
parts = line.split()
try:
positions.append([float(v) for v in parts[1:4]])
except ValueError:
raise ImportFailed(f"{where}, atom {number}: {line.strip()!r}") from None
if len(parts) < 4:
raise ImportFailed(f"{where}, atom {number}: no x, y and z")
symbols.append(_element(parts[0], f"{where}, atom {number}"))
comment = lines[row + 1] if row + 1 < len(lines) else ""
stated = {
"charge" if key.lower() == "charge" else "multiplicity": int(value)
for key, value in _XYZ_STATED.findall(comment)
}
frames.append(
Structure(
id=where,
symbols=symbols,
positions=positions,
charge=stated.get("charge"),
multiplicity=stated.get("multiplicity"),
)
)
row += 2 + count
if not frames:
raise ImportFailed("the XYZ file holds no structure")
from vibesketch.limits import MAX_LAYOUT_ATOMS, LimitExceeded, check_atoms
check_atoms(sum(len(f.symbols) for f in frames), "the XYZ file")
largest = max(len(f.symbols) for f in frames)
if largest > MAX_LAYOUT_ATOMS:
raise LimitExceeded(
f"an XYZ frame of {largest:,} atoms is more than the {MAX_LAYOUT_ATOMS:,} "
"vibe-sketch lays out in 2D"
)
try:
return from_structures(frames, source="xyz")
except NothingDrawn as error:
raise ImportFailed(str(error)) from None
def _calculation(data: str | bytes, fmt: str) -> Derivation:
"""An ORCA output or Hessian file as a document, its graph inferred and recorded so (#86)."""
from vibesketch.document.coordinates import NothingDrawn, from_structures
from vibesketch.document.io import with_edit
from vibesketch.document.orca import OrcaUnreadable, hessian_structures, output_structures
from vibesketch.limits import MAX_LAYOUT_ATOMS, LimitExceeded, check_atoms
reader = output_structures if fmt == "orca" else hessian_structures
try:
structures, warnings = reader(data)
except OrcaUnreadable as error:
raise ImportFailed(str(error)) from None
check_atoms(sum(len(s.symbols) for s in structures), f"the {fmt} file")
if max(len(s.symbols) for s in structures) > MAX_LAYOUT_ATOMS:
raise LimitExceeded(
f"a geometry of more than {MAX_LAYOUT_ATOMS:,} atoms is more than vibe-sketch lays out"
)
try:
found = from_structures(structures, warnings=warnings, source="orca", bond_orders=True)
except NothingDrawn as error:
raise ImportFailed(str(error)) from None
what = "an ORCA output" if fmt == "orca" else "an ORCA Hessian file"
recorded = with_edit(
found.document,
actor="system",
client="import",
summary=f"Opened {what}: bonds and bond orders inferred from its coordinates",
)
recorded = recorded.model_copy(update={"revision": recorded.revision + 1})
return Derivation(recorded, found.warnings, found.result)
def _with_conformers(derived: Derivation, found: Derivation) -> Derivation:
"""``derived`` with the 3D structures ``found`` has for its species that have none."""
from pydantic import ValidationError
from vibesketch.document.model import Species
given = {s.inchikey or s.smiles or s.id: s for s in found.document.species if s.conformers}
species = []
for entry in derived.document.species:
source = given.get(entry.inchikey or entry.smiles or entry.id)
if source is not None and not entry.conformers:
update = {"conformers": source.conformers, "multiplicity": source.multiplicity}
with contextlib.suppress(ValidationError): # the drawing's charge differs
entry = Species.model_validate(entry.model_dump(by_alias=True) | update)
species.append(entry)
document = derived.document.model_copy(update={"species": species})
return Derivation(
Document.model_validate(document.to_json_data()), [*found.warnings, *derived.warnings]
)
def _whole(data: str | bytes, fmt: str) -> Derivation:
"""The whole document a ``.vsk`` or QVF holds; raises :class:`ImportFailed`.
The layer is derived again from the drawing, which keeps ids, labels,
names, conformers and jobs, so it cannot disagree with what is drawn; a
document without a drawing keeps its species as stored.
"""
from vibesketch.document.io import DocumentError, loads
if fmt == "vsk":
try:
stored = Derivation(loads(_text(data)))
except (DocumentError, UnicodeDecodeError) as error:
raise ImportFailed(str(error)) from None
else:
from vibesketch.document.qvf import QvfError, from_qvf
if isinstance(data, str):
raise ImportFailed("a QVF is binary; send its bytes (base64 over REST)")
try:
stored = from_qvf(data)
except QvfError as error:
raise ImportFailed(str(error)) from None
if not stored.document.canvas.data:
return stored
derived = derive(stored.document)
return Derivation(derived.document, [*stored.warnings, *derived.warnings], stored.result)
def _with_definitions(ket: dict[str, Any], definitions: list[Any]) -> None:
"""Draw R-group ``definitions`` in ``ket`` to the right of what it holds (#90)."""
from vibesketch.document.ket import GAP, rgroup_node
xs = [
float(atom["location"][0])
for name, node in ket.items()
if isinstance(node, dict) and node.get("type") == "molecule"
for atom in node.get("atoms", [])
]
x = (max(xs) if xs else 0.0) + 2 * GAP
for definition in definitions:
logic: dict[str, Any] = {}
if definition.occurrence:
logic["range"] = definition.occurrence
if definition.rest_h:
logic["resth"] = True
if definition.if_then:
logic["ifthen"] = definition.if_then
node, width = rgroup_node(definition.number, logic, definition.members, (x, 0.0))
name = f"rg{definition.number}"
ket[name] = node
ket["root"]["nodes"].append({"$ref": name})
x += width + 2 * GAP
def _multiple_groups(molecules: list[Any]) -> list[str]:
"""A note for each multiple group (MUL) the drawing does not carry over (#90)."""
from rdkit import Chem
from vibesketch.document.multiple_groups import collapse_multiple_groups
# A CTfile MUL without its parent subset cannot be verified. Do not
# mistake it for the contracted KET group produced by a verified collapse,
# or infer extra atoms from its label. Preserve the file graph and warn.
missing_parents = 0
for index, mol in enumerate(molecules):
groups = Chem.GetMolSubstanceGroups(mol)
retained = [g for g in groups if g.GetProp("TYPE") != "MUL" or g.GetParentAtoms()]
missing_parents += len(groups) - len(retained)
if len(retained) != len(groups):
cleaned = Chem.Mol(mol)
Chem.ClearMolSubstanceGroups(cleaned)
for group in retained:
Chem.AddMolSubstanceGroup(cleaned, group)
molecules[index] = cleaned
molecules[:] = [collapse_multiple_groups(m) for m in molecules]
count = missing_parents + sum(
1
for mol in molecules
for group in Chem.GetMolSubstanceGroups(mol)
if group.GetProp("TYPE") == "MUL" and group.GetParentAtoms()
)
if not count:
return []
return [
f"{count} multiple group{'s' if count > 1 else ''} (MUL) in the file "
f"{'are' if count > 1 else 'is'} drawn as the file lists the atoms, every copy, "
"without the bracket"
]
[docs]
def read_structures(data: str | bytes, fmt: str) -> Imported:
"""Read structures in ``fmt``; raises :class:`ImportFailed`."""
from rdkit import Chem as chem_module
from rdkit import rdBase
from rdkit.Chem import rdChemReactions
from vibesketch.chem._rdkit import untyped
Chem = untyped(chem_module)
blocker = rdBase.BlockLogs()
try:
if fmt == "smiles":
# Let RDKit read CXSMILES before normalization, so relative groups
# are not discarded as if they were merely an optional name.
lines = [line.strip() for line in _text(data).splitlines() if line.strip()]
if not lines:
raise ImportFailed("no SMILES given")
for line in lines:
_check_cx_stereo(line)
if len(lines) == 1 and ">" in lines[0]:
rxn = rdChemReactions.ReactionFromSmarts(lines[0], useSmiles=True)
if rxn is None:
raise ImportFailed(f"RDKit could not read the reaction SMILES {lines[0]!r}")
return Imported(reaction=_reaction(rxn, "reaction SMILES"))
return Imported(
molecules=[
_sanitized(Chem.MolFromSmiles(s, sanitize=False), f"SMILES {s!r}")
for s in lines
]
)
if fmt == "mol":
from vibesketch.document.rgfile import RgfileUnreadable, is_rgfile, read_rgfile
if is_rgfile(_text(data)):
try:
root, definitions = read_rgfile(_text(data))
except RgfileUnreadable as error:
raise ImportFailed(f"could not read the R-group file: {error}") from None
for definition in definitions:
definition.members = [
_sanitized(m, f"member of R{definition.number}") for m in definition.members
]
molecules = [_sanitized(root, "R-group file's core")]
return Imported(
molecules=molecules,
warnings=_multiple_groups(molecules),
rgroups=definitions,
)
molecules = [
_sanitized(
Chem.MolFromMolBlock(_text(data), sanitize=False, removeHs=False), "molfile"
)
]
return Imported(molecules=molecules, warnings=_multiple_groups(molecules))
if fmt == "sdf":
raw = data if isinstance(data, bytes) else data.encode("utf-8")
supplier = Chem.ForwardSDMolSupplier(io.BytesIO(raw), sanitize=False, removeHs=False)
molecules = [_sanitized(m, "SD record") for m in supplier]
if not molecules:
raise ImportFailed("the SD file holds no records")
return Imported(molecules=molecules, warnings=_multiple_groups(molecules))
if fmt == "rxn":
rxn = rdChemReactions.ReactionFromRxnBlock(_text(data), sanitize=False, removeHs=False)
if rxn is None:
raise ImportFailed("RDKit could not read the RXN file")
return Imported(reaction=_reaction(rxn, "RXN file"))
if fmt in ("cdxml", "cdx"):
from vibesketch.document.cdx import CdxUnreadable, cdx_to_cdxml
from vibesketch.document.cdxml import CdxmlUnreadable, read_cdxml
what = "CDXML" if fmt == "cdxml" else "CDX"
try:
if fmt == "cdx":
if isinstance(data, str):
raise ImportFailed(
"a CDX file is binary; send its bytes (base64 over REST)"
)
data = cdx_to_cdxml(data)
drawing, warnings = read_cdxml(data)
except (CdxmlUnreadable, CdxUnreadable) as error:
raise ImportFailed(f"could not read the {what} file: {error}") from None
return Imported(ket=drawing, warnings=warnings)
if fmt == "inchi":
lines = [line.strip() for line in _text(data).splitlines() if line.strip()]
return Imported(
molecules=[
_sanitized(Chem.MolFromInchi(s.split()[0]), f"InChI {s!r}") for s in lines
]
)
if fmt == "name":
return _named(_text(data).strip())
if fmt in ("peptide", "dna", "rna", "helm"):
from vibesketch.document.sequences import (
SequenceRefused,
helm_molecules,
sequence_molecules,
)
try:
built = (
helm_molecules(_text(data))
if fmt == "helm"
else sequence_molecules(_text(data), fmt) # type: ignore[arg-type]
)
except SequenceRefused as error:
raise ImportFailed(str(error)) from None
return Imported(
molecules=[_sanitized(m, f"{fmt} structure") for m in built], sequenced=True
)
if fmt == "qvf":
return Imported(ket=_whole(data, fmt).document.canvas.data)
if fmt == "xyz":
return Imported(ket=_xyz(data).document.canvas.data)
if fmt in ("ket", "vsk"):
from vibesketch.limits import check_depth
parsed = _parsed_json(_text(data), f"the {fmt} data")
check_depth(parsed, f"the {fmt} data")
ket = parsed["canvas"]["data"] if fmt == "vsk" else parsed
if not isinstance(ket, dict):
raise ImportFailed("not a KET document")
return Imported(ket=ket)
except ImportFailed:
raise
except ValueError as error:
from vibesketch.limits import LimitExceeded
if isinstance(error, LimitExceeded):
raise ImportFailed(str(error)) from None
raise ImportFailed(f"could not read the {fmt} data: {error}") from error
except (KeyError, UnicodeDecodeError) as error:
raise ImportFailed(f"could not read the {fmt} data: {error}") from error
finally:
del blocker
raise ImportFailed(f"unknown import format {fmt!r}; known: {', '.join(IMPORT_FORMATS)}")
# ── Placing imports on the canvas ───────────────────────────────────────
def _points(ket: dict[str, Any]) -> list[tuple[float, float]]:
points: list[tuple[float, float]] = []
for node in ket.get("root", {}).get("nodes", []):
if "$ref" in node:
for atom in ket.get(node["$ref"], {}).get("atoms", []):
points.append((float(atom["location"][0]), float(atom["location"][1])))
elif node.get("type") in ("arrow", "simpleObject"):
points.extend((float(p["x"]), float(p["y"])) for p in node["data"]["pos"])
elif node.get("type") == "plus":
points.append((float(node["location"][0]), float(node["location"][1])))
elif node.get("type") == "text" and "boundingBox" in node:
box = node["boundingBox"]
points.append((float(box["x"]), float(box["y"])))
points.append((float(box["x"]) + float(box["width"]), float(box["y"])))
elif node.get("type") == "image" and "boundingBox" in node:
box = node["boundingBox"]
points.append((float(box["x"]), float(box["y"])))
points.append(
(float(box["x"]) + float(box["width"]), float(box["y"]) - float(box["height"]))
)
return points
def _shifted(ket: dict[str, Any], dx: float, dy: float) -> dict[str, Any]:
import copy
moved = copy.deepcopy(ket)
for node in moved.get("root", {}).get("nodes", []):
if "$ref" in node:
for atom in moved.get(node["$ref"], {}).get("atoms", []):
atom["location"][0] += dx
atom["location"][1] += dy
elif node.get("type") in ("arrow", "simpleObject"):
for point in node["data"]["pos"]:
point["x"] += dx
point["y"] += dy
elif node.get("type") == "plus":
node["location"][0] += dx
node["location"][1] += dy
elif node.get("type") in ("text", "image") and "boundingBox" in node:
node["boundingBox"]["x"] += dx
node["boundingBox"]["y"] += dy
return moved
def _merged(base: dict[str, Any], addition: dict[str, Any], *, gap: float = 3.0) -> dict[str, Any]:
"""``addition`` placed to the right of what ``base`` holds."""
if not base or not base.get("root", {}).get("nodes"):
return addition
old, new = _points(base), _points(addition)
if old and new:
dx = max(p[0] for p in old) + gap - min(p[0] for p in new)
dy = (min(p[1] for p in old) + max(p[1] for p in old)) / 2 - (
min(p[1] for p in new) + max(p[1] for p in new)
) / 2
addition = _shifted(addition, dx, dy)
merged: dict[str, Any] = {k: v for k, v in base.items()}
merged["root"] = dict(base["root"])
nodes = list(base["root"].get("nodes", []))
taken = {k for k in base if k.startswith("mol")}
counter = 0
for node in addition["root"]["nodes"]:
if "$ref" in node:
while f"mol{counter}" in taken:
counter += 1
name = f"mol{counter}"
taken.add(name)
merged[name] = addition[node["$ref"]]
nodes.append({"$ref": name})
else:
nodes.append(node)
merged["root"]["nodes"] = nodes
return merged
def _parsed_json(text: str, what: str) -> Any:
"""JSON, with nesting too deep for the parser itself refused as a limit."""
import json
from vibesketch.limits import MAX_JSON_DEPTH, LimitExceeded
try:
return json.loads(text)
except RecursionError:
raise LimitExceeded(f"{what} nests more than {MAX_JSON_DEPTH} levels deep") from None
def _atoms_in(ket: dict[str, Any]) -> int:
return sum(
len(molecule.get("atoms", []))
for key, molecule in ket.items()
if key.startswith("mol") and isinstance(molecule, dict)
)
def _check_layout(molecules: list[Any]) -> None:
"""Refuse structures too large to lay out, before the layout is attempted (#69)."""
from vibesketch.limits import MAX_LAYOUT_ATOMS, MAX_LAYOUT_TOTAL, LimitExceeded, check_atoms
check_atoms(sum(m.GetNumAtoms() for m in molecules), "the imported data")
undrawn = [
mol for mol in molecules if not (mol.GetNumConformers() and not mol.GetConformer().Is3D())
]
for mol in undrawn:
if mol.GetNumAtoms() > MAX_LAYOUT_ATOMS:
raise LimitExceeded(
f"a structure of {mol.GetNumAtoms():,} atoms has no 2D coordinates, and "
f"vibe-sketch lays out at most {MAX_LAYOUT_ATOMS:,}; import it as MOL or KET "
"with coordinates"
)
total = sum(mol.GetNumAtoms() for mol in undrawn)
if total > MAX_LAYOUT_TOTAL:
raise LimitExceeded(
f"the {len(undrawn):,} structures without 2D coordinates have {total:,} atoms, "
f"and vibe-sketch lays out at most {MAX_LAYOUT_TOTAL:,} in one import; import "
"fewer at a time, or as MOL, SDF or KET with coordinates"
)
[docs]
def import_into(
document: Document | None,
data: str | bytes,
fmt: str,
*,
mode: Literal["append", "replace"] = "append",
) -> Derivation:
"""Read ``data`` and draw it on ``document``'s canvas (a new document if ``None``).
Input beyond :mod:`vibesketch.limits` is refused with :class:`ImportFailed`.
"""
from pydantic import ValidationError
from vibesketch.document.io import validation_errors
from vibesketch.limits import LimitExceeded, check_size
try:
if fmt != "qvf": # an archive is bounded by what it unpacks to
check_size(data, f"the {fmt} data")
return _import(document, data, fmt, mode)
except LimitExceeded as error:
raise ImportFailed(str(error)) from None
except ValidationError as error:
reasons = "; ".join(validation_errors(error)[:3])
raise ImportFailed(f"the drawing is not a valid canvas: {reasons}") from None
def _import(
document: Document | None, data: str | bytes, fmt: str, mode: Literal["append", "replace"]
) -> Derivation:
from vibesketch.limits import check_atoms
if fmt in WHOLE_DOCUMENTS and (document is None or mode == "replace"):
# The whole document: species, conformers, labels and jobs, not just the drawing.
whole = _whole(data, fmt)
if document is None:
return whole
kept = {"canvas", "species", "reactions", "jobs", "style"}
update = {name: getattr(whole.document, name) for name in kept}
replaced = document.model_copy(update=update)
return Derivation(
Document.model_validate(replaced.to_json_data()), whole.warnings, whole.result
)
if fmt in ("xyz", *CALCULATIONS): # the drawing, and the coordinates as 3D structures
found = _xyz(data) if fmt == "xyz" else _calculation(data, fmt)
if document is None:
return found
ket = found.document.canvas.data
if mode == "append":
ket = _merged(document.canvas.data, ket)
drawn = document.model_copy(update={"canvas": Canvas(format="ket", data=ket)})
return _with_conformers(derive(drawn), found)
imported = read_structures(data, fmt)
if imported.ket is not None:
ket = imported.ket
check_atoms(_atoms_in(ket), "the imported drawing")
try:
read_scene(ket)
except (TypeError, AttributeError, KeyError, IndexError, ValueError) as error:
raise ImportFailed(f"the drawing is not a valid canvas: {error}") from None
elif imported.reaction is not None:
_check_layout([m for side in imported.reaction for m in side])
ket = write_scene([], reaction=imported.reaction, conditions=imported.conditions)
else:
_check_layout(imported.molecules)
ket = write_scene(imported.molecules)
if imported.rgroups:
_with_definitions(ket, imported.rgroups)
base = document if document is not None else new_document()
if mode == "append":
ket = _merged(base.canvas.data, ket)
updated = base.model_copy(update={"canvas": Canvas(format="ket", data=ket)})
derived = derive(updated)
if imported.sequenced:
from vibesketch.document.sequences import with_sequences
recorded = with_sequences(derived.document, imported.molecules)
derived = Derivation(recorded, derived.warnings, derived.result)
if not imported.warnings:
return derived
return Derivation(derived.document, [*imported.warnings, *derived.warnings], derived.result)
# ── Export ──────────────────────────────────────────────────────────────
[docs]
def species_mols(document: Document) -> dict[str, Any]:
"""An RDKit molecule per species: with canvas coordinates when drawn."""
from vibesketch.chem.layout import compute_2d
from vibesketch.chem.species import identity, mol_from_smiles, without_hydrogens
by_key: dict[str, Any] = {}
if document.canvas.data:
for fragment in read_scene(document.canvas.data).fragments:
molecule = document.canvas.data.get(fragment.node)
if isinstance(molecule, dict) and (reason := stereo_reason(molecule)):
raise ExportRefused(reason)
if fragment.mol is not None:
by_key.setdefault(identity(fragment.mol).key, fragment.mol)
result: dict[str, Any] = {}
for species in document.species:
key = species.inchikey or species.smiles or species.id
mol = by_key.get(key)
if mol is None and species.smiles:
mol = mol_from_smiles(species.smiles)
compute_2d(mol)
if mol is not None:
result[species.id] = without_hydrogens(mol)
return result
def _with_superatoms(document: Document, mols: dict[str, Any]) -> dict[str, Any]:
"""``mols`` with each drawn species' contracted abbreviations as SUP S-groups.
The species is matched to its first drawn fragment as :func:`species_mols`
matches it, so the structure written is the same.
"""
from vibesketch.chem.species import identity, without_hydrogens
from vibesketch.document.abbreviations import with_superatoms
by_key: dict[str, Any] = {}
if document.canvas.data:
for fragment in read_scene(document.canvas.data).fragments:
if fragment.mol is not None:
by_key.setdefault(identity(fragment.mol).key, fragment)
result = dict(mols)
for species in document.species:
drawn = by_key.get(species.inchikey or species.smiles or species.id)
if species.id in result and drawn is not None and drawn.abbreviations:
result[species.id] = without_hydrogens(with_superatoms(drawn.mol, drawn.abbreviations))
return result
#: The canvas's copolymer subtypes as a molfile writes them (MDL ``SST``).
_MOLFILE_SUBTYPES = {"alt": "ALT", "ran": "RAN", "block": "BLO"}
def _rgroup_parts(document: Document) -> tuple[list[Any], list[Any]]:
"""The R-group cores and definitions to write as an R-group file (#90).
Each core is the molecule figures draw, its R-group atoms numbered for
the molfile; each definition's members carry their attachment points.
"""
from rdkit import Chem
from vibesketch.document.generics import cxsmiles
from vibesketch.document.rgfile import Definition
cores: list[Any] = []
definitions: list[Definition] = []
wanted = {g.cxsmiles for g in document.generics if g.kind == "rgroup-core"}
logic = {r.number: r for r in document.rgroups}
if not document.canvas.data or not (wanted or logic):
return cores, definitions
written: set[str] = set()
for fragment in read_scene(document.canvas.data).fragments:
if fragment.mol is not None or fragment.figure is None:
continue
node = document.canvas.data.get(fragment.node)
if isinstance(node, dict) and node.get("type") == "rgroup":
rlogic = node.get("rlogic")
value = rlogic.get("number") if isinstance(rlogic, dict) else None
number = value if isinstance(value, int) else 0
if number not in logic:
continue
members = Chem.RWMol(fragment.figure)
Chem.CollapseAttachmentPoints(members)
entry = logic[number]
definitions.append(
Definition(
number,
list(Chem.GetMolFrags(members.GetMol(), asMols=True)),
entry.occurrence or "",
entry.rest_h,
entry.if_then or 0,
)
)
continue
text = cxsmiles(fragment.figure)
if text not in wanted or text in written:
continue
written.add(text)
core = Chem.Mol(fragment.figure)
for atom in core.GetAtoms():
label = atom.GetProp("dummyLabel") if atom.HasProp("dummyLabel") else ""
if atom.GetAtomicNum() == 0 and label[1:].isdigit():
atom.SetIntProp("_MolFileRLabel", int(label[1:]))
cores.append(core)
return cores, sorted(definitions, key=lambda d: d.number)
def _bracketed_mols(document: Document) -> dict[str, Any]:
"""Each bracketed or labelled generic structure's molecule, in molfile terms (#90, #139).
A labelled attachment atom is written as a pseudo-atom with its label as
the alias (``A`` with an ``A`` line in V2000).
The molecule is the one figures draw, with its brackets. A molfile names
a copolymer's kind in ``SST`` and the connectivity in upper case; a
multiple group (MUL) expands every copy with the parent atom subset
(MDL CTfile formats, S-group properties), which the canvas draws once.
"""
from rdkit import Chem
from vibesketch.document.generics import cxsmiles
wanted = {g.cxsmiles: g.id for g in document.generics if g.kind in ("bracketed", "labelled")}
found: dict[str, Any] = {}
if not wanted or not document.canvas.data:
return found
for fragment in read_scene(document.canvas.data).fragments:
if fragment.mol is not None or fragment.figure is None:
continue
identifier = wanted.get(cxsmiles(fragment.figure))
if identifier is None or identifier in found:
continue
mol = Chem.Mol(fragment.figure)
for atom in mol.GetAtoms():
if atom.HasProp("attachmentLabel"):
label = str(atom.GetProp("attachmentLabel"))
for name in list(atom.GetPropNames()):
atom.ClearProp(name)
# MDL's pseudo-atom: an "A" atom whose alias is the label (CTfile, alias lines).
atom.SetProp("dummyLabel", "A")
atom.SetProp("molFileAlias", label)
from vibesketch.document.multiple_groups import MultipleGroupRefused, expand_multiple_groups
try:
mol = expand_multiple_groups(mol)
except MultipleGroupRefused as error:
raise ExportRefused(str(error)) from None
for index in range(len(Chem.GetMolSubstanceGroups(mol))):
# The group itself, not a copy: its properties are rewritten in place.
group = Chem.GetMolSubstanceGroupWithIdx(mol, index)
kind = group.GetProp("TYPE")
if group.HasProp("CONNECT"):
group.SetProp("CONNECT", str(group.GetProp("CONNECT")).upper())
if kind == "COP":
subtype = _MOLFILE_SUBTYPES.get(str(group.GetProp("LABEL")), "")
group.ClearProp("LABEL")
if subtype:
group.SetProp("SUBTYPE", subtype)
found[identifier] = mol
return found
def _molfile_molecule(mol: Any) -> Any:
"""Copy a molecule with a flag consistent with its absolute tetrahedral centers."""
from rdkit import Chem
copy = Chem.Mol(mol)
relative = {
atom.GetIdx()
for group in copy.GetStereoGroups()
if group.GetGroupType() != Chem.StereoGroupType.STEREO_ABSOLUTE
for atom in group.GetAtoms()
}
# A zero molfile chiral flag makes ungrouped wedge centers relative to
# readers honoring it. Enhanced OR/AND collections remain explicit in V3000.
# https://www.rdkit.org/docs/source/rdkit.Chem.rdmolops.html#rdkit.Chem.rdmolops.TranslateChiralFlagToStereoGroups
absolute = any(
atom.GetChiralTag()
in (Chem.ChiralType.CHI_TETRAHEDRAL_CW, Chem.ChiralType.CHI_TETRAHEDRAL_CCW)
and atom.GetIdx() not in relative
for atom in copy.GetAtoms()
)
copy.SetIntProp("_MolFileChiralFlag", int(absolute))
return copy
def _named(name: str) -> Imported:
"""The structure for ``name``: vibe-qc's naming engine first, then OPSIN (#94).
The structure's source is said in a warning, so the user sees which route named it.
"""
from rdkit import Chem
from vibesketch.naming import NamingUnavailable, smiles_for_name
from vibesketch.naming import install_hint as engine_hint
from vibesketch.naming.opsin import OpsinUnavailable, read_name
tried: list[str] = []
try:
smiles = smiles_for_name(name)
except NamingUnavailable:
smiles = None
tried.append(engine_hint())
else:
if smiles is None:
tried.append("the naming engine has no structure for it")
source = "vibe-qc's naming engine"
if smiles is None:
try:
found = read_name(name)
except OpsinUnavailable as error:
tried.append(str(error))
else:
if found.smiles is None:
tried.append(f"OPSIN: {found.reason}")
smiles, source = found.smiles, "OPSIN"
if smiles is None:
raise ImportFailed(f"no structure for {name!r}: " + "; ".join(tried))
mol = _sanitized(Chem.MolFromSmiles(smiles, sanitize=False), "named structure")
return Imported(molecules=[mol], warnings=[f"{name!r} read by {source}"])
[docs]
def export(
document: Document,
fmt: str,
*,
style: str | None = None,
dpi: int = 300,
stereo_labels: bool = False,
sequences: bool = False,
) -> Payload:
"""``document`` in ``fmt``; raises :class:`ExportRefused`.
``stereo_labels`` prints CIP descriptors in SVG, PNG and PDF figures;
``sequences`` prints a species built from a sequence as it (#98).
"""
import json
from rdkit import Chem, rdBase
from rdkit.Chem import inchi as inchi_module
from rdkit.Chem import rdChemReactions
from vibesketch.chem._rdkit import untyped
inchi = untyped(inchi_module)
extension = EXTENSIONS.get(fmt, fmt)
if fmt in ("mol", "sdf", "rxn", "smiles", "smi", "inchi", "cdxml", "helm", "fasta"):
from vibesketch.document.surfaces import SurfaceRefused, records
try:
surfaces = document.surfaces or records(document)
except SurfaceRefused as error:
raise ExportRefused(str(error)) from None
if surfaces:
raise ExportRefused(
f"{fmt} cannot express surface objects; use QVF, VSK, KET or a figure"
)
if fmt == "cdxml":
from vibesketch.document.cdxml import CdxmlRefused, cdxml_source
try:
return Payload(cdxml_source(document, style=style), MEDIA_TYPES[fmt], extension)
except CdxmlRefused as error:
raise ExportRefused(str(error)) from None
if fmt not in EXPORT_FORMATS:
raise ExportRefused(f"unknown export format {fmt!r}; known: {', '.join(EXPORT_FORMATS)}")
if fmt in ("helm", "fasta"):
from vibesketch.document.sequences import SequenceRefused, sequence_export
try:
text = sequence_export(document, fmt) # type: ignore[arg-type]
except SequenceRefused as error:
raise ExportRefused(str(error)) from None
return Payload(text.encode("utf-8"), MEDIA_TYPES[fmt], extension)
if fmt == "qvf":
from vibesketch.document.qvf import QvfError, to_qvf
try:
return Payload(to_qvf(document), MEDIA_TYPES[fmt], extension)
except QvfError as error:
raise ExportRefused(str(error)) from None
if fmt == "vsk":
return Payload(dumps(document).encode("utf-8"), MEDIA_TYPES[fmt], extension)
if fmt == "ket":
text = json.dumps(document.canvas.data, indent=2, ensure_ascii=False) + "\n"
return Payload(text.encode("utf-8"), MEDIA_TYPES[fmt], extension)
if fmt == "chemfig":
from vibesketch.render.chemfig import ChemfigRefused, chemfig_source
try:
return Payload(chemfig_source(document).encode("utf-8"), MEDIA_TYPES[fmt], extension)
except ChemfigRefused as error:
raise ExportRefused(str(error)) from None
if fmt in ("svg", "png", "pdf", "eps", "emf"):
from vibesketch.render import render
content = render(
document,
fmt,
style=style if style is not None else document.style.preset,
dpi=dpi,
stereo_labels=stereo_labels,
sequences=sequences,
)
return Payload(content, MEDIA_TYPES[fmt], extension)
# Bracketed groups, such as polymers, and structures bonded to labelled
# attachment atoms travel in MOL and SDF beside the species,
# and R-group cores with their definitions as an R-group file in MOL (#90).
bracketed = _bracketed_mols(document) if fmt in ("mol", "sdf") else {}
cores, definitions = _rgroup_parts(document) if fmt in ("mol", "sdf") else ([], [])
if fmt == "sdf" and (cores or definitions):
raise ExportRefused(
"an SDF file cannot hold R-group definitions; export MOL, which writes them "
"as an R-group file"
)
solids = [s for s in document.species if s.solid is not None]
if solids and fmt in ("mol", "sdf", "rxn", "smiles", "inchi"):
names = ", ".join(s.solid.text for s in solids if s.solid is not None)
raise ExportRefused(
f"{fmt.upper()} cannot write a solid given by its formula ({names}); export the "
"drawing as KET, CDXML, .vsk or a figure"
)
mols = species_mols(document)
if not mols and not bracketed and not cores:
raise ExportRefused("the document has no species to export")
if fmt in ("mol", "sdf", "rxn"):
# Contracted abbreviations travel as superatoms; the structure is the same.
mols = _with_superatoms(document, mols)
if fmt == "mol":
combined = None
for mol in [*mols.values(), *bracketed.values(), *cores]:
combined = mol if combined is None else Chem.CombineMols(combined, mol)
if cores or definitions:
from vibesketch.document.rgfile import rgfile
text = rgfile(_molfile_molecule(combined), definitions)
return Payload(text.encode("utf-8"), MEDIA_TYPES[fmt], extension)
return Payload(
str(Chem.MolToMolBlock(_molfile_molecule(combined))).encode("utf-8"),
MEDIA_TYPES[fmt],
extension,
)
if fmt == "sdf":
stream = io.StringIO()
writer = Chem.SDWriter(stream)
for species in document.species:
if species.id not in mols:
continue
record = _molfile_molecule(mols[species.id])
record.SetProp("_Name", species.label or species.id)
record.SetProp("vibesketch_id", species.id)
if species.name is not None:
record.SetProp("name", species.name.value)
record.SetProp("charge", str(species.charge))
record.SetProp("multiplicity", str(species.multiplicity))
writer.write(record)
for generic in document.generics:
if generic.id not in bracketed:
continue
record = _molfile_molecule(bracketed[generic.id])
record.SetProp("_Name", generic.label or generic.id)
record.SetProp("vibesketch_id", generic.id)
writer.write(record)
writer.close()
return Payload(stream.getvalue().encode("utf-8"), MEDIA_TYPES[fmt], extension)
if fmt == "rxn":
if not document.reactions:
raise ExportRefused("an RXN file needs a reaction, and the document has none")
if len(document.reactions) > 1:
raise ExportRefused(
f"an RXN file holds one reaction, and the document has {len(document.reactions)}"
)
reaction = document.reactions[0]
rxn = rdChemReactions.ChemicalReaction()
for sid in reaction.reactants:
rxn.AddReactantTemplate(_molfile_molecule(mols[sid]))
for sid in reaction.reagents:
rxn.AddAgentTemplate(_molfile_molecule(mols[sid]))
for sid in reaction.products:
rxn.AddProductTemplate(_molfile_molecule(mols[sid]))
block = str(rdChemReactions.ReactionToRxnBlock(rxn, separateAgents=True))
return Payload(block.encode("utf-8"), MEDIA_TYPES[fmt], extension)
if fmt == "smiles":
lines: list[str] = []
for reaction in document.reactions:
def side(ids: list[str]) -> str:
return ".".join(str(Chem.MolToSmiles(mols[i])) for i in ids)
lines.append(
f"{side(reaction.reactants)}>{side(reaction.reagents)}>{side(reaction.products)}"
f" {reaction.id}"
)
if not lines:
lines = [
f"{Chem.MolToSmiles(mols[s.id])} {s.label or s.id}"
for s in document.species
if s.id in mols
]
return Payload(("\n".join(lines) + "\n").encode("utf-8"), MEDIA_TYPES[fmt], extension)
# inchi
rows = []
blocker_lines = [(s, mols[s.id]) for s in document.species if s.id in mols]
for species, mol in blocker_lines:
blocker = rdBase.BlockLogs() # "Proton(s) added/removed" for H+
text = inchi.MolToInchi(mol)
del blocker
if not text:
raise ExportRefused(f"InChI cannot represent species {species.id}")
rows.append(f"{text} {species.label or species.id}")
return Payload(("\n".join(rows) + "\n").encode("utf-8"), MEDIA_TYPES[fmt], extension)