Refactor disorder export demo for Sonar

Author: SchrodingersCattt

scripts/materials/disorder_workflow_demo/demo_csd_disorder_workflow.py

@@ -42,6 +42,10 @@
 LICENSE_REQUIREMENT = "CSD-Materials or another licence tier with CSD Python API database access"
 
 
+def _escape_cif_string(value: str | None) -> str:
+    return (value or "?").replace("'", "''")
+
+
 def _format_float(value: Any, precision: int = 4) -> str:
     if value is None:
         return "?"
@@ -62,94 +66,99 @@ def _format_esd(value: Any, esd: Any, precision: int = 4) -> str:
     return base
 
 
-def export_full_cif_from_csd(refcode: str, output_path: Path) -> dict[str, Any]:
-    """Export a CSD structure to CIF with occupancy and disorder metadata."""
-    from ccdc.io import EntryReader
+def _build_disorder_map(crystal: Any) -> dict[str, tuple[int | None, str | None]]:
+    disorder_map: dict[str, tuple[int | None, str | None]] = {}
+    disorder = getattr(crystal, "disorder", None)
+    if disorder is None:
+        return disorder_map
 
-    with EntryReader("CSD") as reader:
-        entry = reader.entry(refcode)
-        crystal = reader.crystal(refcode)
+    try:
+        for assembly in disorder.assemblies:
+            assembly_id = str(getattr(assembly, "id", ""))
+            for group in assembly.groups:
+                group_id = int(getattr(group, "id", 0))
+                for atom in group.atoms:
+                    disorder_map[str(atom.label)] = (group_id, assembly_id)
+    except Exception:
+        return disorder_map
+    return disorder_map
 
-    molecule = crystal.disordered_molecule or crystal.molecule
 
-    disorder_map: dict[str, tuple[int | None, str | None]] = {}
-    disorder = getattr(crystal, "disorder", None)
-    if disorder is not None:
-        try:
-            for assembly in disorder.assemblies:
-                assembly_id = str(getattr(assembly, "id", ""))
-                for group in assembly.groups:
-                    group_id = int(getattr(group, "id", 0))
-                    for atom in group.atoms:
-                        disorder_map[str(atom.label)] = (group_id, assembly_id)
-        except Exception:
-            pass
+def _extract_occupancy(atom: Any) -> float:
+    try:
+        if atom.occupancy is not None:
+            return float(atom.occupancy)
+    except Exception:
+        pass
+    return 1.0
 
+
+def _extract_displacement(atom: Any) -> tuple[str, float | None, dict[str, Any] | None]:
+    displacement_type = "?"
+    u_iso = None
+    aniso = None
+    try:
+        displacement = atom.displacement_parameters
+        if displacement is None:
+            return displacement_type, u_iso, aniso
+        u_iso = float(displacement.isotropic_equivalent)
+        if displacement.type != "Anisotropic":
+            return "Uiso", u_iso, aniso
+        values = displacement.values
+        uncertainties = displacement.uncertainties
+        aniso = {
+            "u11": values[0][0], "u22": values[1][1], "u33": values[2][2],
+            "u12": values[0][1], "u13": values[0][2], "u23": values[1][2],
+            "e11": uncertainties[0][0], "e22": uncertainties[1][1], "e33": uncertainties[2][2],
+            "e12": uncertainties[0][1], "e13": uncertainties[0][2], "e23": uncertainties[1][2],
+        }
+        return "Uani", u_iso, aniso
+    except Exception:
+        return displacement_type, u_iso, aniso
+
+
+def _extract_atom_data(molecule: Any, disorder_map: dict[str, tuple[int | None, str | None]]) -> list[dict[str, Any]]:
     atom_data: list[dict[str, Any]] = []
     for atom in molecule.atoms:
         label = str(atom.label) if atom.label else "?"
         symbol = str(atom.atomic_symbol) if atom.atomic_symbol else "?"
         fractional = atom.fractional_coordinates
-        occupancy = 1.0
-        try:
-            if atom.occupancy is not None:
-                occupancy = float(atom.occupancy)
-        except Exception:
-            pass
-
-        displacement_type = "?"
-        u_iso = None
-        aniso = None
-        try:
-            displacement = atom.displacement_parameters
-            if displacement is not None:
-                u_iso = float(displacement.isotropic_equivalent)
-                if displacement.type == "Anisotropic":
-                    displacement_type = "Uani"
-                    values = displacement.values
-                    uncertainties = displacement.uncertainties
-                    aniso = {
-                        "u11": values[0][0], "u22": values[1][1], "u33": values[2][2],
-                        "u12": values[0][1], "u13": values[0][2], "u23": values[1][2],
-                        "e11": uncertainties[0][0], "e22": uncertainties[1][1], "e33": uncertainties[2][2],
-                        "e12": uncertainties[0][1], "e13": uncertainties[0][2], "e23": uncertainties[1][2],
-                    }
-                else:
-                    displacement_type = "Uiso"
-        except Exception:
-            pass
-
+        displacement_type, u_iso, aniso = _extract_displacement(atom)
         disorder_group, disorder_assembly = disorder_map.get(label, (None, None))
         atom_data.append({
             "label": label,
             "symbol": symbol,
             "fx": float(fractional.x) if fractional else None,
             "fy": float(fractional.y) if fractional else None,
             "fz": float(fractional.z) if fractional else None,
-            "occupancy": occupancy,
+            "occupancy": _extract_occupancy(atom),
             "u_iso": u_iso,
             "displacement_type": displacement_type,
             "disorder_group": disorder_group,
             "disorder_assembly": disorder_assembly,
             "aniso": aniso,
         })
+    return atom_data
+
 
+def _extract_bonds(molecule: Any) -> list[tuple[str, str, float, str]]:
     bonds: list[tuple[str, str, float, str]] = []
     try:
         for bond in molecule.bonds:
             atom_1, atom_2 = bond.atoms
             bonds.append((str(atom_1.label), str(atom_2.label), float(bond.length), str(bond.bond_type)))
     except Exception:
-        pass
+        return bonds
+    return bonds
+
 
+def _build_header_lines(refcode: str, entry: Any, crystal: Any) -> list[str]:
     cell_a, cell_b, cell_c = crystal.cell_lengths
     alpha, beta, gamma = crystal.cell_angles
-    disorder_present = any(atom["disorder_group"] is not None for atom in atom_data)
-
-    lines = [
+    return [
         f"data_{refcode}",
         "_audit_creation_method            'CSD Python API disorder workflow demo'",
-        f"_chemical_name_common             '{(entry.chemical_name or '?').replace(chr(39), chr(39) + chr(39))}'",
+        f"_chemical_name_common             '{_escape_cif_string(entry.chemical_name)}'",
         f"_chemical_formula_sum             '{entry.formula or '?'}'",
         f"_cell_length_a                    {_format_float(cell_a)}",
         f"_cell_length_b                    {_format_float(cell_b)}",
@@ -163,13 +172,20 @@ def export_full_cif_from_csd(refcode: str, output_path: Path) -> dict[str, Any]:
         "",
     ]
 
-    if crystal.symmetry_operators:
-        lines.extend(["loop_", "_symmetry_equiv_pos_as_xyz"])
-        for operator in crystal.symmetry_operators:
-            lines.append(f"  '{operator}'")
-        lines.append("")
 
-    lines.extend([
+def _build_symmetry_lines(crystal: Any) -> list[str]:
+    if not crystal.symmetry_operators:
+        return []
+    lines = ["loop_", "_symmetry_equiv_pos_as_xyz"]
+    for operator in crystal.symmetry_operators:
+        lines.append(f"  '{operator}'")
+    lines.append("")
+    return lines
+
+
+def _build_atom_site_lines(atom_data: list[dict[str, Any]]) -> list[str]:
+    disorder_present = any(atom["disorder_group"] is not None for atom in atom_data)
+    lines = [
         "loop_",
         "_atom_site_label",
         "_atom_site_type_symbol",
@@ -179,7 +195,7 @@ def export_full_cif_from_csd(refcode: str, output_path: Path) -> dict[str, Any]:
         "_atom_site_occupancy",
         "_atom_site_U_iso_or_equiv",
         "_atom_site_thermal_displace_type",
-    ])
+    ]
     if disorder_present:
         lines.extend([
             "_atom_site_disorder_assembly",
@@ -200,52 +216,96 @@ def export_full_cif_from_csd(refcode: str, output_path: Path) -> dict[str, Any]:
             f"{atom['displacement_type']:>5s}",
         ]
         if disorder_present:
-            row.append(f"{str(atom['disorder_assembly']) if atom['disorder_assembly'] is not None else '.':>4s}")
-            row.append(f"{str(atom['disorder_group']) if atom['disorder_group'] is not None else '.':>4s}")
+            row.extend([
+                f"{str(atom['disorder_assembly']) if atom['disorder_assembly'] is not None else '.':>4s}",
+                f"{str(atom['disorder_group']) if atom['disorder_group'] is not None else '.':>4s}",
+            ])
         lines.append("  " + " ".join(row))
     lines.append("")
+    return lines
+
 
+def _build_anisotropic_lines(atom_data: list[dict[str, Any]]) -> list[str]:
     anisotropic_atoms = [atom for atom in atom_data if atom["aniso"]]
-    if anisotropic_atoms:
-        lines.extend([
-            "loop_",
-            "_atom_site_aniso_label",
-            "_atom_site_aniso_U_11",
-            "_atom_site_aniso_U_22",
-            "_atom_site_aniso_U_33",
-            "_atom_site_aniso_U_12",
-            "_atom_site_aniso_U_13",
-            "_atom_site_aniso_U_23",
-        ])
-        for atom in anisotropic_atoms:
-            u = atom["aniso"]
-            lines.append(
-                f"  {atom['label']:<8s}"
-                f" {_format_esd(u['u11'], u['e11']):>12s}"
-                f" {_format_esd(u['u22'], u['e22']):>12s}"
-                f" {_format_esd(u['u33'], u['e33']):>12s}"
-                f" {_format_esd(u['u12'], u['e12']):>12s}"
-                f" {_format_esd(u['u13'], u['e13']):>12s}"
-                f" {_format_esd(u['u23'], u['e23']):>12s}"
-            )
-        lines.append("")
-
-    if bonds:
-        lines.extend([
-            "loop_",
-            "_geom_bond_atom_site_label_1",
-            "_geom_bond_atom_site_label_2",
-            "_geom_bond_distance",
-            "_ccdc_geom_bond_type",
-        ])
-        for atom_1, atom_2, distance, bond_type in bonds:
-            lines.append(f"  {atom_1:<8s} {atom_2:<8s} {distance:8.4f} {bond_type}")
-        lines.append("")
+    if not anisotropic_atoms:
+        return []
+    lines = [
+        "loop_",
+        "_atom_site_aniso_label",
+        "_atom_site_aniso_U_11",
+        "_atom_site_aniso_U_22",
+        "_atom_site_aniso_U_33",
+        "_atom_site_aniso_U_12",
+        "_atom_site_aniso_U_13",
+        "_atom_site_aniso_U_23",
+    ]
+    for atom in anisotropic_atoms:
+        u = atom["aniso"]
+        lines.append(
+            f"  {atom['label']:<8s}"
+            f" {_format_esd(u['u11'], u['e11']):>12s}"
+            f" {_format_esd(u['u22'], u['e22']):>12s}"
+            f" {_format_esd(u['u33'], u['e33']):>12s}"
+            f" {_format_esd(u['u12'], u['e12']):>12s}"
+            f" {_format_esd(u['u13'], u['e13']):>12s}"
+            f" {_format_esd(u['u23'], u['e23']):>12s}"
+        )
+    lines.append("")
+    return lines
 
-    lines.append("#END")
+
+def _build_bond_lines(bonds: list[tuple[str, str, float, str]]) -> list[str]:
+    if not bonds:
+        return []
+    lines = [
+        "loop_",
+        "_geom_bond_atom_site_label_1",
+        "_geom_bond_atom_site_label_2",
+        "_geom_bond_distance",
+        "_ccdc_geom_bond_type",
+    ]
+    for atom_1, atom_2, distance, bond_type in bonds:
+        lines.append(f"  {atom_1:<8s} {atom_2:<8s} {distance:8.4f} {bond_type}")
+    lines.append("")
+    return lines
+
+
+def _write_cif_file(output_path: Path, lines: list[str]) -> None:
     output_path.parent.mkdir(parents=True, exist_ok=True)
     output_path.write_text("\n".join(lines), encoding="utf-8")
 
+
+def _validated_output_path(output_dir: Path, refcode: str) -> Path:
+    base_dir = output_dir.expanduser().resolve()
+    candidate = (base_dir / f"{refcode}.cif").resolve()
+    try:
+        candidate.relative_to(base_dir)
+    except ValueError as exc:
+        raise ValueError(f"Refusing to write outside output directory: {candidate}") from exc
+    return candidate
+
+
+def export_full_cif_from_csd(refcode: str, output_path: Path) -> dict[str, Any]:
+    """Export a CSD structure to CIF with occupancy and disorder metadata."""
+    from ccdc.io import EntryReader
+
+    with EntryReader("CSD") as reader:
+        entry = reader.entry(refcode)
+        crystal = reader.crystal(refcode)
+
+    molecule = crystal.disordered_molecule or crystal.molecule
+    disorder_map = _build_disorder_map(crystal)
+    atom_data = _extract_atom_data(molecule, disorder_map)
+    bonds = _extract_bonds(molecule)
+    lines = []
+    lines.extend(_build_header_lines(refcode, entry, crystal))
+    lines.extend(_build_symmetry_lines(crystal))
+    lines.extend(_build_atom_site_lines(atom_data))
+    lines.extend(_build_anisotropic_lines(atom_data))
+    lines.extend(_build_bond_lines(bonds))
+    lines.append("#END")
+    _write_cif_file(output_path, lines)
+
     return {
         "refcode": refcode,
         "output": str(output_path),
@@ -268,7 +328,7 @@ def main() -> None:
 
     print(f"Licence requirement: {LICENSE_REQUIREMENT}")
     for refcode in args.refcodes:
-        output_path = output_dir / f"{refcode}.cif"
+        output_path = _validated_output_path(output_dir, refcode)
         summary = export_full_cif_from_csd(refcode, output_path)
         print(
             f"Exported {summary['refcode']} -> {summary['output']} "