From a0e006af22568841830598ecb4cd9adffe76d956 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Thu, 10 Sep 2026 16:40:45 +0100 Subject: [PATCH 1/5] auto-generate marshmallow schema from linkml --- .../fieldextraction/generate_marshmallow.py | 244 ++++ biosim_schema/utils/generate_artifacts.py | 11 + project/biosimdb_fields.py | 1298 +++++++++++++++++ pyproject.toml | 5 + 4 files changed, 1558 insertions(+) create mode 100644 biosim_schema/fieldextraction/generate_marshmallow.py create mode 100644 project/biosimdb_fields.py diff --git a/biosim_schema/fieldextraction/generate_marshmallow.py b/biosim_schema/fieldextraction/generate_marshmallow.py new file mode 100644 index 0000000..7edcace --- /dev/null +++ b/biosim_schema/fieldextraction/generate_marshmallow.py @@ -0,0 +1,244 @@ +"""Generate Invenio-compatible Marshmallow schemas + ES mappings from LinkML. + +LinkML YAML -> SchemaView -> Marshmallow schema classes + FIELDS/MAPPING dicts. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +from linkml_runtime.utils.schemaview import SchemaView + +HEADER = '''"""Generated from BioSimDB LinkML schema. DO NOT EDIT MANUALLY.""" + +from marshmallow import Schema, fields +from marshmallow.validate import OneOf +from marshmallow_utils.fields import SanitizedUnicode + +COMMUNITY = "BioSimDB" +''' + +PRIMITIVE_FIELD = { + "string": "SanitizedUnicode()", + "integer": "fields.Integer()", + "float": "fields.Float()", + "double": "fields.Float()", + "boolean": "fields.Boolean()", + "date": "fields.Date()", + "datetime": "fields.DateTime()", +} + +PRIMITIVE_MAPPING = { + "string": "text", + "integer": "integer", + "float": "double", + "double": "double", + "boolean": "boolean", + "date": "date", + "datetime": "date", +} + + +def schema_class_name(linkml_class_name: str) -> str: + return f"{linkml_class_name}Schema" + + +@dataclass +class ClassBuild: + """Accumulated Marshmallow field lines + ES mapping dict for one class.""" + + name: str + field_lines: list[str] = field(default_factory=list) + mapping: dict = field(default_factory=dict) + depends_on: set[str] = field(default_factory=set) + + +class MarshmallowGenerator: + def __init__(self, schema_path: str | Path): + self.view = SchemaView(str(schema_path)) + self.builds: dict[str, ClassBuild] = {} + + # -- slot -> field/mapping ------------------------------------------------- + + def _field_and_mapping_for_slot(self, cls_name: str, slot_name: str): + slot = self.view.induced_slot(slot_name, cls_name) + range_name = slot.range or "string" + + enum_def = self.view.get_enum(range_name) + class_def = self.view.get_class(range_name) + + if class_def is not None: + nested_name = schema_class_name(range_name) + base_field = f"fields.Nested({nested_name})" + base_mapping = {"type": "object", "properties": {}} # filled in later + depends = {range_name} + elif enum_def is not None: + values = list(enum_def.permissible_values.keys()) + base_field = f"fields.String(validate=OneOf({values!r}))" + base_mapping = {"type": "keyword"} + depends = set() + else: + base_field = PRIMITIVE_FIELD.get(range_name, "SanitizedUnicode()") + base_mapping = {"type": PRIMITIVE_MAPPING.get(range_name, "text")} + depends = set() + + if slot.multivalued: + base_field = f"fields.List({base_field})" + # ES: list of objects/scalars uses the same mapping as a single item + + kwargs = [] + if slot.required: + kwargs.append("required=True") + else: + kwargs.append("allow_none=True") + if kwargs: + if base_field.endswith("()"): + base_field = base_field[:-1] + ", ".join(kwargs) + ")" + else: + base_field = base_field[:-1] + ", " + ", ".join(kwargs) + ")" + + return ( + slot_name, + base_field, + base_mapping, + depends, + class_def is not None, + range_name, + ) + + # -- one class --------------------------------------------------------- + + def _build_class(self, cls_name: str) -> ClassBuild: + build = ClassBuild(name=cls_name) + for slot_name in self.view.class_slots(cls_name): + name, field_expr, mapping, depends, is_nested, range_name = ( + self._field_and_mapping_for_slot(cls_name, slot_name) + ) + build.field_lines.append(f" {name} = {field_expr}") + build.mapping[name] = mapping + build.depends_on |= depends + if is_nested: + # remember which nested class produced this mapping, filled later + build.mapping[name]["_nested_class"] = range_name + return build + + def build_all(self) -> None: + for cls_name in self.view.all_classes(imports=True): + self.builds[cls_name] = self._build_class(cls_name) + self._resolve_nested_mappings() + + def _resolve_nested_mappings(self) -> None: + """Backfill nested ES `properties` once all classes are built (handles recursion).""" + + def resolve(mapping: dict, seen: frozenset[str]) -> dict: + resolved = {} + for key, val in mapping.items(): + if not isinstance(val, dict): + continue + nested_cls = val.pop("_nested_class", None) + if nested_cls: + if nested_cls in seen: + resolved[key] = {"type": "object"} # break recursive cycle + else: + nested_mapping = resolve( + self.builds[nested_cls].mapping, seen | {nested_cls} + ) + resolved[key] = {"type": "object", "properties": nested_mapping} + else: + resolved[key] = val + return resolved + + for build in self.builds.values(): + build.mapping = resolve(build.mapping, frozenset({build.name})) + + # -- ordering + rendering ------------------------------------------------ + + def _topo_order(self) -> list[str]: + visited: set[str] = set() + order: list[str] = [] + + def visit(name: str, stack: tuple[str, ...]): + if name in visited or name not in self.builds: + return + if name in stack: + return # recursive reference, break cycle + visited.add(name) + for dep in self.builds[name].depends_on: + visit(dep, stack + (name,)) + order.append(name) + + for name in self.builds: + visit(name, ()) + return order + + def render_schemas(self) -> str: + parts = [] + for cls_name in self._topo_order(): + build = self.builds[cls_name] + body = "\n".join(build.field_lines) or " pass" + parts.append(f"class {schema_class_name(cls_name)}(Schema):\n{body}\n") + return "\n\n".join(parts) + + def render_root(self, root_class: str | None = None) -> str: + root = root_class or self._find_tree_root() + build = self.builds[root] + + fields_lines = [] + for slot_name in self.view.class_slots(root): + slot = self.view.induced_slot(slot_name, root) + if self.view.get_class(slot.range) is not None: + nested = schema_class_name(slot.range) + fields_lines.append( + f' "{slot_name}": fields.Nested({nested}, allow_none=True),' + ) + else: + _, field_expr, _, _, _, _ = self._field_and_mapping_for_slot( + root, slot_name + ) + fields_lines.append(f' "{slot_name}": {field_expr},') + + fields_block = "FIELDS = {\n" + "\n".join(fields_lines) + "\n}\n" + + import json + + mapping_block = "MAPPING = " + json.dumps(build.mapping, indent=4) + "\n" + mapping_block = mapping_block.replace( + '"type"', '"type"' + ) # keep as valid python dict via json (double quotes are valid py) + + return fields_block + "\n" + mapping_block + + def _find_tree_root(self) -> str: + for cls_name, cls in self.view.all_classes(imports=True).items(): + if getattr(cls, "tree_root", False): + return cls_name + raise ValueError("No tree_root class found in schema") + + def generate(self, root_class: str | None = None) -> str: + self.build_all() + return ( + HEADER + + "\n\n" + + self.render_schemas() + + "\n\n" + + self.render_root(root_class) + ) + + +def main() -> None: + import argparse + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("schema", help="Path to root LinkML schema YAML") + parser.add_argument("-o", "--output", required=True, help="Output .py file") + parser.add_argument("--root-class", default=None) + args = parser.parse_args() + + generator = MarshmallowGenerator(args.schema) + code = generator.generate(root_class=args.root_class) + Path(args.output).write_text(code, encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/biosim_schema/utils/generate_artifacts.py b/biosim_schema/utils/generate_artifacts.py index 630c772..afba803 100644 --- a/biosim_schema/utils/generate_artifacts.py +++ b/biosim_schema/utils/generate_artifacts.py @@ -28,6 +28,7 @@ write_summary_csv, ) from biosim_schema.fieldextraction.extract_webform_fields import WebFormFieldExtractor +from biosim_schema.fieldextraction.generate_marshmallow import MarshmallowGenerator @dataclass(frozen=True) @@ -194,12 +195,22 @@ def generate_summary(paths: ArtifactPaths) -> None: json.dump(summary, fp, indent=2) +def generate_invenio_marshmallow(paths: ArtifactPaths) -> None: + """Generate Invenio-compatible Marshmallow schemas + ES mappings.""" + generator = MarshmallowGenerator(str(paths.schema_path)) + code = generator.generate() + out_path = paths.project_dir / "invenio" / "biosimdb_fields.py" + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(code, encoding="utf-8") + + def generate_derived(paths: ArtifactPaths) -> None: """Generate derived schema artefacts used by downstream tools.""" generate_webform_fields(paths) generate_engine_mappings(paths) generate_summary(paths) + generate_invenio_marshmallow(paths) def fix_generated_docs(linkml_docs_dir: Path) -> None: diff --git a/project/biosimdb_fields.py b/project/biosimdb_fields.py new file mode 100644 index 0000000..d65ea98 --- /dev/null +++ b/project/biosimdb_fields.py @@ -0,0 +1,1298 @@ +"""Generated from BioSimDB LinkML schema. DO NOT EDIT MANUALLY.""" + +from marshmallow import Schema, fields +from marshmallow.validate import OneOf +from marshmallow_utils.fields import SanitizedUnicode + +COMMUNITY = "BioSimDB" + + +class LengthQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['Å', 'nm']), allow_none=True) + + +class VolumeQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['ų', 'nm³']), allow_none=True) + + +class TimeQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['s', 'ms', 'μs', 'ns', 'ps', 'fs']), allow_none=True) + + +class FrequencyQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['1/ps']), allow_none=True) + + +class FrictionCoefficientQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['a/ps', 'kg/s']), allow_none=True) + + +class MolarEnergyQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['kcal/mol', 'kJ/mol']), allow_none=True) + + +class EnergyQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['kWh']), allow_none=True) + + +class TemperatureQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['K', '°C', '°F']), allow_none=True) + + +class PressureQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['bar', 'Pa']), allow_none=True) + + +class CompressibilityQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['1/bar', '1/Pa']), allow_none=True) + + +class MassQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['g/mol', 'Da']), allow_none=True) + + +class ConcentrationQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['M']), allow_none=True) + + +class ForceQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['kJ/mol/nm', 'kcal/mol/Å']), allow_none=True) + + +class ChargeQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['e', 'C']), allow_none=True) + + +class AngleQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['degree', 'radian']), allow_none=True) + + +class ByteQuantitySchema(Schema): + value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['GB', 'MB']), allow_none=True) + + +class VectorLengthQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['Å', 'nm']), allow_none=True) + + +class VectorAngleQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['degree', 'radian']), allow_none=True) + + +class VectorVolumeQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['Å', 'nm']), allow_none=True) + + +class VectorCompressibilityQuantitySchema(Schema): + vector_value = fields.Float(allow_none=True) + value_unit = fields.String(validate=OneOf(['1/bar', '1/Pa']), allow_none=True) + + +class VectorPressureQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['bar', 'Pa']), allow_none=True) + + +class VectorTemperatureQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['K', '°C', '°F']), allow_none=True) + + +class VectorTimeQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['s', 'ms', 'μs', 'ns', 'ps', 'fs']), allow_none=True) + + +class MatrixPressureQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['bar', 'Pa']), allow_none=True) + + +class MatrixCompressibilityQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = fields.String(validate=OneOf(['1/bar', '1/Pa']), allow_none=True) + + +class MatrixQuantitySchema(Schema): + vector_value = fields.List(fields.Float(), allow_none=True) + value_unit = SanitizedUnicode(allow_none=True) + + +class AnalysisSchema(Schema): + analysis_tool = fields.List(fields.String(validate=OneOf(['mdout_analyzer.py', 'ambpdb', 'CPPTRAJ', 'PYTRAJ', 'MMPBSA.py', 'Free Energy Workflow (FEW)', 'edgember', 'SAX-RISM', 'SAX-MD', 'MoFT', 'ndfes', 'PLUMED', 'MDanalysis'])), allow_none=True) + analysis_software = fields.List(fields.String(validate=OneOf(['Visual Molecular Dynamics (VMD)', 'Schrödinger Maestro', 'PyMOL', 'Avogadro'])), allow_none=True) + analysis_method = fields.List(fields.String(validate=OneOf(['RMSD', 'DSSP', 'GIST', 'Hydrogen Bonds', 'Connolly surface', 'Radius of Gyration', 'BAR/PBSA'])), allow_none=True) + + +class SetupSchema(Schema): + setup_tool = fields.String(validate=OneOf(['pdb4amber', 'prepareforleap', 'packmol', 'packmol_memgen', 'LEaP', 'antechamber', 'pyMSMT', 'mdgx', 'parmed']), allow_none=True) + + +class ProductionSchema(Schema): + simulation_tool = fields.List(fields.String(validate=OneOf(['sander', 'pmemd', 'gem.pmemd', 'mdrun'])), allow_none=True) + simulation_software = fields.List(fields.String(validate=OneOf(['Amber', 'GROMACS', 'LAMMPS', 'NAMD', 'OpenMM', 'CHARMM', 'DL_POLY', 'HOOMD-blue', 'Desmond', 'ACEMD', 'CP2K'])), allow_none=True) + simulation_software_version = SanitizedUnicode(allow_none=True) + simulation_method = fields.List(fields.String(validate=OneOf(['Self-guided Langevin Dynamics', 'Accelerated Molecular Dynamics', 'Gaussian Accelerated Molecular Dynamics', 'Targeted Molecular Dynamics', 'Nudged Elastic Band Calculations', 'Adaptive String Method', 'LMOD method', 'DL-FIND Optimization', 'Thermodynamic Integration (TI)', 'Linear Interaction Energies (LIE)', 'Replica Exchange Molecular Dynamics (REMD)', 'Adaptively Biased Molecular Dynamics (ABMD)', 'Steered Molecular Dynamics (SMD)', 'Umbrella Sampling', 'Metadynamics', 'Swarms of Trajectories String Method', 'Constant pH Molecular Dynamics', 'Constant Redox Potential Molecular Dynamics', 'Continuous Constant pH Molecular Dynamics', 'NMR Refinement', 'X-ray and CryoEM Refinement', 'Locally Enhanced Sampling'])), allow_none=True) + + +class EquilibrationSchema(Schema): + simulation_tool = fields.List(fields.String(validate=OneOf(['sander', 'pmemd', 'gem.pmemd', 'mdrun'])), allow_none=True) + simulation_software = fields.List(fields.String(validate=OneOf(['Amber', 'GROMACS', 'LAMMPS', 'NAMD', 'OpenMM', 'CHARMM', 'DL_POLY', 'HOOMD-blue', 'Desmond', 'ACEMD', 'CP2K'])), allow_none=True) + + +class MinimisationSchema(Schema): + energy_tolerance = fields.Nested(ForceQuantitySchema, allow_none=True) + number_of_minimisation_steps = fields.Integer(allow_none=True) + minimisation_distance_step_size = fields.Nested(LengthQuantitySchema, allow_none=True) + minimisation_algorithm = fields.List(fields.String(validate=OneOf(['Steepest Descent', 'Conjugate Gradient', 'L-BFGS', 'XMIN', 'LMOD', 'None'])), allow_none=True) + simulation_tool = fields.List(fields.String(validate=OneOf(['sander', 'pmemd', 'gem.pmemd', 'mdrun'])), allow_none=True) + simulation_software = fields.List(fields.String(validate=OneOf(['Amber', 'GROMACS', 'LAMMPS', 'NAMD', 'OpenMM', 'CHARMM', 'DL_POLY', 'HOOMD-blue', 'Desmond', 'ACEMD', 'CP2K'])), allow_none=True) + + +class SimulationStagesSchema(Schema): + setup = fields.Nested(SetupSchema, allow_none=True) + minimisation = fields.Nested(MinimisationSchema, allow_none=True) + equilibration = fields.Nested(EquilibrationSchema, allow_none=True) + production = fields.Nested(ProductionSchema, allow_none=True) + analysis = fields.Nested(AnalysisSchema, allow_none=True) + + +class IntegratorSchema(Schema): + integrator_algorithm = fields.String(validate=OneOf(['Velocity-Verlet', 'Leap-frog', 'Verlet', 'Euler']), allow_none=True) + frame_step = fields.Nested(TimeQuantitySchema, allow_none=True) + time_step = fields.Nested(TimeQuantitySchema, allow_none=True) + number_of_steps = fields.Integer(allow_none=True) + simulation_time = fields.Nested(TimeQuantitySchema, allow_none=True) + + +class BarostatSchema(Schema): + barostat_algorithm = fields.String(validate=OneOf(['Berendsen', 'Andersen', 'Parrinello-Rahman', 'Nose-Hoover', 'Monte Carlo', 'Martyna-Tuckerman-Tobias-Klein']), allow_none=True) + compressibility = fields.Nested(CompressibilityQuantitySchema, allow_none=True) + compressibility_vector = fields.Nested(MatrixCompressibilityQuantitySchema, allow_none=True) + target_pressure = fields.Nested(PressureQuantitySchema, allow_none=True) + target_pressure_vector = fields.Nested(MatrixPressureQuantitySchema, allow_none=True) + pressure_time_constant = fields.Nested(TimeQuantitySchema, allow_none=True) + pressure_coupling_frequency = fields.Nested(FrequencyQuantitySchema, allow_none=True) + pressure_coupling_type = fields.String(validate=OneOf(['isotropic', 'semi-isotropic', 'anisotropic', 'surface tension']), allow_none=True) + + +class ThermostatSchema(Schema): + thermostat_algorithm = fields.String(validate=OneOf(['Langevin', 'Berendsen', 'Andersen', 'Nose-Hoover', 'Bussi']), allow_none=True) + target_temperature = fields.Nested(TemperatureQuantitySchema, allow_none=True) + target_temperature_vector = fields.Nested(VectorTemperatureQuantitySchema, allow_none=True) + collision_frequency = fields.Nested(FrequencyQuantitySchema, allow_none=True) + temperature_time_constant = fields.Nested(VectorTimeQuantitySchema, allow_none=True) + coupling_group = SanitizedUnicode(allow_none=True) + chain_length = fields.Integer(allow_none=True) + friction_coefficient = fields.Nested(FrictionCoefficientQuantitySchema, allow_none=True) + + +class EnsembleSchema(Schema): + ensemble_type = fields.String(validate=OneOf(['NPT', 'NVT', 'NVE', 'μVT']), allow_none=True) + random_seed = fields.Integer(allow_none=True) + + +class InteractionsSchema(Schema): + restraints = fields.Boolean(allow_none=True) + electrostatic_cutoff_distance = fields.Nested(LengthQuantitySchema, allow_none=True) + vdw_cutoff_distance = fields.Nested(LengthQuantitySchema, allow_none=True) + bond_length_constraints_algorithm = fields.String(validate=OneOf(['SHAKE', 'RATTLE', 'SETTLE', 'LINCS', 'CCMA']), allow_none=True) + long_range_interaction_method = fields.String(validate=OneOf(['Cutoff', 'Ewald', 'PME', 'P3M', 'FMM', 'RF']), allow_none=True) + + +class SimulationSettingsSchema(Schema): + ensemble = fields.Nested(EnsembleSchema, allow_none=True) + integrator = fields.Nested(IntegratorSchema, allow_none=True) + barostat = fields.Nested(BarostatSchema, allow_none=True) + thermostat = fields.Nested(ThermostatSchema, allow_none=True) + interactions = fields.Nested(InteractionsSchema, allow_none=True) + + +class MoleculeIDSchema(Schema): + PDB_ID = SanitizedUnicode(allow_none=True) + UNIPROT_ID = SanitizedUnicode(allow_none=True) + SMILES = SanitizedUnicode(allow_none=True) + InChI = SanitizedUnicode(allow_none=True) + InChIKey = SanitizedUnicode(allow_none=True) + alphafold_ID = SanitizedUnicode(allow_none=True) + PubChem_CID = SanitizedUnicode(allow_none=True) + protein_sequence = SanitizedUnicode(allow_none=True) + nucleic_sequence = SanitizedUnicode(allow_none=True) + predicted_structure = fields.Boolean(allow_none=True) + modified = fields.Boolean(allow_none=True) + molecular_formula = SanitizedUnicode(allow_none=True) + molecular_weight = fields.Nested(MassQuantitySchema, allow_none=True) + molecule_charge = fields.Nested(ChargeQuantitySchema, allow_none=True) + molecule_count = fields.Integer(allow_none=True) + atom_count = fields.Integer(allow_none=True) + monomer_count = fields.Integer(allow_none=True) + simulated_particle_names = SanitizedUnicode(allow_none=True) + simulated_molecule_name = SanitizedUnicode(allow_none=True) + + +class SystemCountsSchema(Schema): + total_molecule_count = fields.Integer(allow_none=True) + total_atom_count = fields.Integer(allow_none=True) + unique_molecule_count = fields.Integer(allow_none=True) + salt_concentration = fields.Nested(ConcentrationQuantitySchema, allow_none=True) + + +class SystemCompositionSchema(Schema): + system_counts = fields.Nested(SystemCountsSchema, allow_none=True) + molecule_ID = fields.List(fields.Nested(MoleculeIDSchema), allow_none=True) + + +class SimulationAveragesSchema(Schema): + average_kinetic_energy = fields.Nested(MolarEnergyQuantitySchema, allow_none=True) + average_potential_energy = fields.Nested(MolarEnergyQuantitySchema, allow_none=True) + average_enthalpy = fields.Nested(MolarEnergyQuantitySchema, allow_none=True) + average_pressure = fields.Nested(PressureQuantitySchema, allow_none=True) + average_temperature = fields.Nested(TemperatureQuantitySchema, allow_none=True) + average_volume = fields.Nested(VolumeQuantitySchema, allow_none=True) + average_volume_vector = fields.Nested(VectorVolumeQuantitySchema, allow_none=True) + + +class SimulationObservablesSchema(Schema): + simulation_averages = fields.Nested(SimulationAveragesSchema, allow_none=True) + + +class ParticlesSchema(Schema): + masses = fields.Boolean(allow_none=True) + fixed_charges = fields.Boolean(allow_none=True) + system_charge = fields.Nested(ChargeQuantitySchema, allow_none=True) + coarse_grained = fields.Boolean(allow_none=True) + resolution = fields.String(validate=OneOf(['All Atom', 'United Atom', 'Coarse-Grained', 'Mesoscale']), allow_none=True) + + +class ConnectivitySchema(Schema): + bonds = fields.Boolean(allow_none=True) + dihedrals = fields.Boolean(allow_none=True) + + +class TopologyMetadataSchema(Schema): + connectivity = fields.Nested(ConnectivitySchema, allow_none=True) + particles = fields.Nested(ParticlesSchema, allow_none=True) + + +class SimulationBoxSchema(Schema): + box_dimensions = fields.Nested(VectorLengthQuantitySchema, allow_none=True) + box_angles = fields.Nested(VectorAngleQuantitySchema, allow_none=True) + box_type = fields.String(validate=OneOf(['Cubic', 'Tetragonal', 'Orthorhombic', 'Truncated Octahedron', 'Triclinic']), allow_none=True) + periodic_boundary_conditions = fields.String(validate=OneOf(['None', 'xyz', 'xy', 'xz', 'yz']), allow_none=True) + + +class TrajectoriesSchema(Schema): + positions = fields.Boolean(allow_none=True) + forces = fields.Boolean(allow_none=True) + velocities = fields.Boolean(allow_none=True) + polarizable_charges = fields.Boolean(allow_none=True) + energies = fields.Boolean(allow_none=True) + water = fields.Boolean(allow_none=True) + replica = fields.Boolean(allow_none=True) + frame_count = fields.Integer(allow_none=True) + + +class TrajectoryMetadataSchema(Schema): + simulation_box = fields.Nested(SimulationBoxSchema, allow_none=True) + trajectory_output = fields.Nested(TrajectoriesSchema, allow_none=True) + + +class NucleicPotentialSchema(Schema): + nucleic_potential_name = fields.String(validate=OneOf(['ff99-bsc0', 'ff99OL3', 'LJbb', 'ROC', 'Shaw', 'OL15', 'OL21', 'OL24', 'OL3', 'bsc1', 'terminal_monophosphate']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class WaterPotentialSchema(Schema): + water_potential_name = fields.String(validate=OneOf(['OPC', 'OPC3', 'OPC3POL', 'POL3', 'TIP3P', 'TIP3PFB', 'TIP4PFB', 'TIP4P', 'TIP5P', 'TIP4PEW', 'SPCE', 'SPCEB', 'SPC/Fw', 'q-SPC/Fw']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class LipidPotentialSchema(Schema): + lipid_potential_name = fields.String(validate=OneOf(['LIPID21']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class ProteinPotentialSchema(Schema): + protein_potential_name = fields.String(validate=OneOf(['ff19SB', 'ff99SB', 'ff99SB-ILDN', 'ff99SB-disp', 'ff14SB', 'ff14SBonlysc', 'ff15ipq', 'fb15', 'ff03', 'ff03ua', 'phosaa10', 'phosaa14SB', 'phosaa19SB', 'ff14SB_modAA', 'ff19SB_modAA']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class CarbohydratePotentialSchema(Schema): + carbohydrate_potential_name = fields.String(validate=OneOf(['GLYCAM06', 'GLYCAM_06EP', 'GLYCAM_06j-1']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class PolymerPotentialSchema(Schema): + polymer_potential_name = fields.String(validate=OneOf(['LignAmb25']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class GeneralPotentialSchema(Schema): + general_potential_name = fields.String(validate=OneOf(['gem.pmemd', 'GAFF', 'GAFF2', 'OPLS', 'GROMOS', 'CHARMM']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class MachineLearnedPotentialSchema(Schema): + machine_learned_potential_name = fields.String(validate=OneOf(['MACE', 'ANI', 'NequIP', 'UMA', 'AceFF']), allow_none=True) + modified = fields.Boolean(allow_none=True) + + +class PotentialMetadataSchema(Schema): + water_potential = fields.Nested(WaterPotentialSchema, allow_none=True) + protein_potential = fields.Nested(ProteinPotentialSchema, allow_none=True) + lipid_potential = fields.Nested(LipidPotentialSchema, allow_none=True) + nucleic_potential = fields.Nested(NucleicPotentialSchema, allow_none=True) + carbohydrate_potential = fields.Nested(CarbohydratePotentialSchema, allow_none=True) + polymer_potential = fields.Nested(PolymerPotentialSchema, allow_none=True) + general_potential = fields.Nested(GeneralPotentialSchema, allow_none=True) + machine_learned_potential = fields.Nested(MachineLearnedPotentialSchema, allow_none=True) + + +class SoftwareSchema(Schema): + operating_system = fields.String(validate=OneOf(['Linux', 'macOS', 'Windows']), allow_none=True) + scheduler = fields.String(validate=OneOf(['SLURM', 'PBS', 'LSF', 'SGE', 'None']), allow_none=True) + MPI_library = fields.String(validate=OneOf(['OpenMPI', 'MPICH', 'IntelMPI', 'MVAPICH2', 'None']), allow_none=True) + container_runtime = fields.String(validate=OneOf(['Apptainer', 'Docker', 'Podman', 'None']), allow_none=True) + + +class PerformanceSchema(Schema): + wall_time = fields.Nested(TimeQuantitySchema, allow_none=True) + energy_consumption = fields.Nested(EnergyQuantitySchema, allow_none=True) + + +class HardwareSchema(Schema): + execution_platform = fields.String(validate=OneOf(['HPC Cluster', 'Cloud VM', 'Local']), allow_none=True) + node_type = fields.String(validate=OneOf(['CPU only', 'GPU Accelerated', 'Hybrid CPU GPU']), allow_none=True) + node_count = fields.Integer(allow_none=True) + CPU_vendor = fields.String(validate=OneOf(['AMD', 'Intel', 'ARM', 'Other']), allow_none=True) + CPU_architecture = fields.String(validate=OneOf(['x86', 'ARM']), allow_none=True) + sockets_per_node = fields.Integer(allow_none=True) + cores_per_socket = fields.Integer(allow_none=True) + threads_per_core = fields.Integer(allow_none=True) + GPU_vendor = fields.String(validate=OneOf(['Nvidia', 'AMD', 'Intel', 'None']), allow_none=True) + GPUs_per_node = fields.Integer(allow_none=True) + memory_per_node = fields.Nested(ByteQuantitySchema, allow_none=True) + + +class ComputationalEnvironmentSchema(Schema): + hardware = fields.Nested(HardwareSchema, allow_none=True) + software = fields.Nested(SoftwareSchema, allow_none=True) + performance = fields.Nested(PerformanceSchema, allow_none=True) + + +class FileMetadataSchema(Schema): + file_name = SanitizedUnicode(required=True) + file_size = fields.Nested(ByteQuantitySchema, allow_none=True) + file_hash = SanitizedUnicode(allow_none=True) + file_hash_algorithm = fields.String(validate=OneOf(['sha256', 'md5', 'xxh3_64']), allow_none=True) + file_role = fields.String(validate=OneOf(['topology', 'trajectory', 'metadata', 'log', 'parameter', 'other']), allow_none=True) + + +class SimulationMetadataSchema(Schema): + stages = fields.Nested(SimulationStagesSchema, allow_none=True) + settings = fields.Nested(SimulationSettingsSchema, allow_none=True) + observables = fields.Nested(SimulationObservablesSchema, allow_none=True) + topology = fields.Nested(TopologyMetadataSchema, allow_none=True) + trajectory = fields.Nested(TrajectoryMetadataSchema, allow_none=True) + composition = fields.Nested(SystemCompositionSchema, allow_none=True) + potentials = fields.Nested(PotentialMetadataSchema, allow_none=True) + compute = fields.Nested(ComputationalEnvironmentSchema, allow_none=True) + files = fields.List(fields.Nested(FileMetadataSchema), allow_none=True) + biosim_schema_version = SanitizedUnicode(allow_none=True) + + +FIELDS = { + "stages": fields.Nested(SimulationStagesSchema, allow_none=True), + "settings": fields.Nested(SimulationSettingsSchema, allow_none=True), + "observables": fields.Nested(SimulationObservablesSchema, allow_none=True), + "topology": fields.Nested(TopologyMetadataSchema, allow_none=True), + "trajectory": fields.Nested(TrajectoryMetadataSchema, allow_none=True), + "composition": fields.Nested(SystemCompositionSchema, allow_none=True), + "potentials": fields.Nested(PotentialMetadataSchema, allow_none=True), + "compute": fields.Nested(ComputationalEnvironmentSchema, allow_none=True), + "files": fields.Nested(FileMetadataSchema, allow_none=True), + "biosim_schema_version": SanitizedUnicode(allow_none=True), +} + +MAPPING = { + "stages": { + "type": "object", + "properties": { + "setup": { + "type": "object", + "properties": { + "setup_tool": { + "type": "keyword" + } + } + }, + "minimisation": { + "type": "object", + "properties": { + "energy_tolerance": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "number_of_minimisation_steps": { + "type": "integer" + }, + "minimisation_distance_step_size": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "minimisation_algorithm": { + "type": "keyword" + }, + "simulation_tool": { + "type": "keyword" + }, + "simulation_software": { + "type": "keyword" + } + } + }, + "equilibration": { + "type": "object", + "properties": { + "simulation_tool": { + "type": "keyword" + }, + "simulation_software": { + "type": "keyword" + } + } + }, + "production": { + "type": "object", + "properties": { + "simulation_tool": { + "type": "keyword" + }, + "simulation_software": { + "type": "keyword" + }, + "simulation_software_version": { + "type": "text" + }, + "simulation_method": { + "type": "keyword" + } + } + }, + "analysis": { + "type": "object", + "properties": { + "analysis_tool": { + "type": "keyword" + }, + "analysis_software": { + "type": "keyword" + }, + "analysis_method": { + "type": "keyword" + } + } + } + } + }, + "settings": { + "type": "object", + "properties": { + "ensemble": { + "type": "object", + "properties": { + "ensemble_type": { + "type": "keyword" + }, + "random_seed": { + "type": "integer" + } + } + }, + "integrator": { + "type": "object", + "properties": { + "integrator_algorithm": { + "type": "keyword" + }, + "frame_step": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "time_step": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "number_of_steps": { + "type": "integer" + }, + "simulation_time": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + } + } + }, + "barostat": { + "type": "object", + "properties": { + "barostat_algorithm": { + "type": "keyword" + }, + "compressibility": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "compressibility_vector": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "target_pressure": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "target_pressure_vector": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "pressure_time_constant": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "pressure_coupling_frequency": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "pressure_coupling_type": { + "type": "keyword" + } + } + }, + "thermostat": { + "type": "object", + "properties": { + "thermostat_algorithm": { + "type": "keyword" + }, + "target_temperature": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "target_temperature_vector": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "collision_frequency": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "temperature_time_constant": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "coupling_group": { + "type": "text" + }, + "chain_length": { + "type": "integer" + }, + "friction_coefficient": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + } + } + }, + "interactions": { + "type": "object", + "properties": { + "restraints": { + "type": "boolean" + }, + "electrostatic_cutoff_distance": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "vdw_cutoff_distance": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "bond_length_constraints_algorithm": { + "type": "keyword" + }, + "long_range_interaction_method": { + "type": "keyword" + } + } + } + } + }, + "observables": { + "type": "object", + "properties": { + "simulation_averages": { + "type": "object", + "properties": { + "average_kinetic_energy": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "average_potential_energy": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "average_enthalpy": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "average_pressure": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "average_temperature": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "average_volume": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "average_volume_vector": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + } + } + } + } + }, + "topology": { + "type": "object", + "properties": { + "connectivity": { + "type": "object", + "properties": { + "bonds": { + "type": "boolean" + }, + "dihedrals": { + "type": "boolean" + } + } + }, + "particles": { + "type": "object", + "properties": { + "masses": { + "type": "boolean" + }, + "fixed_charges": { + "type": "boolean" + }, + "system_charge": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "coarse_grained": { + "type": "boolean" + }, + "resolution": { + "type": "keyword" + } + } + } + } + }, + "trajectory": { + "type": "object", + "properties": { + "simulation_box": { + "type": "object", + "properties": { + "box_dimensions": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "box_angles": { + "type": "object", + "properties": { + "vector_value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "box_type": { + "type": "keyword" + }, + "periodic_boundary_conditions": { + "type": "keyword" + } + } + }, + "trajectory_output": { + "type": "object", + "properties": { + "positions": { + "type": "boolean" + }, + "forces": { + "type": "boolean" + }, + "velocities": { + "type": "boolean" + }, + "polarizable_charges": { + "type": "boolean" + }, + "energies": { + "type": "boolean" + }, + "water": { + "type": "boolean" + }, + "replica": { + "type": "boolean" + }, + "frame_count": { + "type": "integer" + } + } + } + } + }, + "composition": { + "type": "object", + "properties": { + "system_counts": { + "type": "object", + "properties": { + "total_molecule_count": { + "type": "integer" + }, + "total_atom_count": { + "type": "integer" + }, + "unique_molecule_count": { + "type": "integer" + }, + "salt_concentration": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + } + } + }, + "molecule_ID": { + "type": "object", + "properties": { + "PDB_ID": { + "type": "text" + }, + "UNIPROT_ID": { + "type": "text" + }, + "SMILES": { + "type": "text" + }, + "InChI": { + "type": "text" + }, + "InChIKey": { + "type": "text" + }, + "alphafold_ID": { + "type": "text" + }, + "PubChem_CID": { + "type": "text" + }, + "protein_sequence": { + "type": "text" + }, + "nucleic_sequence": { + "type": "text" + }, + "predicted_structure": { + "type": "boolean" + }, + "modified": { + "type": "boolean" + }, + "molecular_formula": { + "type": "text" + }, + "molecular_weight": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "molecule_charge": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "molecule_count": { + "type": "integer" + }, + "atom_count": { + "type": "integer" + }, + "monomer_count": { + "type": "integer" + }, + "simulated_particle_names": { + "type": "text" + }, + "simulated_molecule_name": { + "type": "text" + } + } + } + } + }, + "potentials": { + "type": "object", + "properties": { + "water_potential": { + "type": "object", + "properties": { + "water_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "protein_potential": { + "type": "object", + "properties": { + "protein_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "lipid_potential": { + "type": "object", + "properties": { + "lipid_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "nucleic_potential": { + "type": "object", + "properties": { + "nucleic_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "carbohydrate_potential": { + "type": "object", + "properties": { + "carbohydrate_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "polymer_potential": { + "type": "object", + "properties": { + "polymer_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "general_potential": { + "type": "object", + "properties": { + "general_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + }, + "machine_learned_potential": { + "type": "object", + "properties": { + "machine_learned_potential_name": { + "type": "keyword" + }, + "modified": { + "type": "boolean" + } + } + } + } + }, + "compute": { + "type": "object", + "properties": { + "hardware": { + "type": "object", + "properties": { + "execution_platform": { + "type": "keyword" + }, + "node_type": { + "type": "keyword" + }, + "node_count": { + "type": "integer" + }, + "CPU_vendor": { + "type": "keyword" + }, + "CPU_architecture": { + "type": "keyword" + }, + "sockets_per_node": { + "type": "integer" + }, + "cores_per_socket": { + "type": "integer" + }, + "threads_per_core": { + "type": "integer" + }, + "GPU_vendor": { + "type": "keyword" + }, + "GPUs_per_node": { + "type": "integer" + }, + "memory_per_node": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + } + } + }, + "software": { + "type": "object", + "properties": { + "operating_system": { + "type": "keyword" + }, + "scheduler": { + "type": "keyword" + }, + "MPI_library": { + "type": "keyword" + }, + "container_runtime": { + "type": "keyword" + } + } + }, + "performance": { + "type": "object", + "properties": { + "wall_time": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "energy_consumption": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + } + } + } + } + }, + "files": { + "type": "object", + "properties": { + "file_name": { + "type": "text" + }, + "file_size": { + "type": "object", + "properties": { + "value": { + "type": "double" + }, + "value_unit": { + "type": "keyword" + } + } + }, + "file_hash": { + "type": "text" + }, + "file_hash_algorithm": { + "type": "keyword" + }, + "file_role": { + "type": "keyword" + } + } + }, + "biosim_schema_version": { + "type": "text" + } +} diff --git a/pyproject.toml b/pyproject.toml index caf8d46..f2e64ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,11 @@ Source = "https://github.com/jkalayan/biosim-schema" # [project.optional-dependencies] +invenio = [ + "marshmallow>=3.20,<4", + "marshmallow-utils>=0.9,<1", +] + testing = [ "pytest>=9.0,<10.0", "pytest-cov>=7.0,<8.0", From 70e1aa1ea860956e0ed4df27bddb5a4dc833c832 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Fri, 11 Sep 2026 11:40:24 +0100 Subject: [PATCH 2/5] include invenio schema in artifact gen --- .../fieldextraction/generate_marshmallow.py | 6 +- biosim_schema/utils/generate_artifacts.py | 5 + .../{ => linkml/invenio}/biosimdb_fields.py | 98 +++++++++---------- 3 files changed, 58 insertions(+), 51 deletions(-) rename project/{ => linkml/invenio}/biosimdb_fields.py (100%) diff --git a/biosim_schema/fieldextraction/generate_marshmallow.py b/biosim_schema/fieldextraction/generate_marshmallow.py index 7edcace..4298b8f 100644 --- a/biosim_schema/fieldextraction/generate_marshmallow.py +++ b/biosim_schema/fieldextraction/generate_marshmallow.py @@ -51,7 +51,7 @@ class ClassBuild: name: str field_lines: list[str] = field(default_factory=list) mapping: dict = field(default_factory=dict) - depends_on: set[str] = field(default_factory=set) + depends_on: list[str] = field(default_factory=list) class MarshmallowGenerator: @@ -117,7 +117,9 @@ def _build_class(self, cls_name: str) -> ClassBuild: ) build.field_lines.append(f" {name} = {field_expr}") build.mapping[name] = mapping - build.depends_on |= depends + for dep in depends: + if dep not in build.depends_on: + build.depends_on.append(dep) if is_nested: # remember which nested class produced this mapping, filled later build.mapping[name]["_nested_class"] = range_name diff --git a/biosim_schema/utils/generate_artifacts.py b/biosim_schema/utils/generate_artifacts.py index afba803..8d614c9 100644 --- a/biosim_schema/utils/generate_artifacts.py +++ b/biosim_schema/utils/generate_artifacts.py @@ -54,6 +54,10 @@ def jsonld_path(self) -> Path: def jsonschema_path(self) -> Path: return self.project_dir / "jsonschema" / f"{self.schema_name}.schema.json" + @property + def invenio_path(self) -> Path: + return self.project_dir / "invenio" / f"{self.schema_name}_fields.py" + @property def docs_dir(self) -> Path: return self.repo_root / "docs" @@ -371,6 +375,7 @@ def main() -> None: paths.summary_yaml_path, paths.summary_csv_path, paths.jsonschema_path, + paths.invenio_path, ] if args.include_docs: diff --git a/project/biosimdb_fields.py b/project/linkml/invenio/biosimdb_fields.py similarity index 100% rename from project/biosimdb_fields.py rename to project/linkml/invenio/biosimdb_fields.py index d65ea98..283d7eb 100644 --- a/project/biosimdb_fields.py +++ b/project/linkml/invenio/biosimdb_fields.py @@ -137,21 +137,17 @@ class MatrixQuantitySchema(Schema): value_unit = SanitizedUnicode(allow_none=True) -class AnalysisSchema(Schema): - analysis_tool = fields.List(fields.String(validate=OneOf(['mdout_analyzer.py', 'ambpdb', 'CPPTRAJ', 'PYTRAJ', 'MMPBSA.py', 'Free Energy Workflow (FEW)', 'edgember', 'SAX-RISM', 'SAX-MD', 'MoFT', 'ndfes', 'PLUMED', 'MDanalysis'])), allow_none=True) - analysis_software = fields.List(fields.String(validate=OneOf(['Visual Molecular Dynamics (VMD)', 'Schrödinger Maestro', 'PyMOL', 'Avogadro'])), allow_none=True) - analysis_method = fields.List(fields.String(validate=OneOf(['RMSD', 'DSSP', 'GIST', 'Hydrogen Bonds', 'Connolly surface', 'Radius of Gyration', 'BAR/PBSA'])), allow_none=True) - - class SetupSchema(Schema): setup_tool = fields.String(validate=OneOf(['pdb4amber', 'prepareforleap', 'packmol', 'packmol_memgen', 'LEaP', 'antechamber', 'pyMSMT', 'mdgx', 'parmed']), allow_none=True) -class ProductionSchema(Schema): +class MinimisationSchema(Schema): + energy_tolerance = fields.Nested(ForceQuantitySchema, allow_none=True) + number_of_minimisation_steps = fields.Integer(allow_none=True) + minimisation_distance_step_size = fields.Nested(LengthQuantitySchema, allow_none=True) + minimisation_algorithm = fields.List(fields.String(validate=OneOf(['Steepest Descent', 'Conjugate Gradient', 'L-BFGS', 'XMIN', 'LMOD', 'None'])), allow_none=True) simulation_tool = fields.List(fields.String(validate=OneOf(['sander', 'pmemd', 'gem.pmemd', 'mdrun'])), allow_none=True) simulation_software = fields.List(fields.String(validate=OneOf(['Amber', 'GROMACS', 'LAMMPS', 'NAMD', 'OpenMM', 'CHARMM', 'DL_POLY', 'HOOMD-blue', 'Desmond', 'ACEMD', 'CP2K'])), allow_none=True) - simulation_software_version = SanitizedUnicode(allow_none=True) - simulation_method = fields.List(fields.String(validate=OneOf(['Self-guided Langevin Dynamics', 'Accelerated Molecular Dynamics', 'Gaussian Accelerated Molecular Dynamics', 'Targeted Molecular Dynamics', 'Nudged Elastic Band Calculations', 'Adaptive String Method', 'LMOD method', 'DL-FIND Optimization', 'Thermodynamic Integration (TI)', 'Linear Interaction Energies (LIE)', 'Replica Exchange Molecular Dynamics (REMD)', 'Adaptively Biased Molecular Dynamics (ABMD)', 'Steered Molecular Dynamics (SMD)', 'Umbrella Sampling', 'Metadynamics', 'Swarms of Trajectories String Method', 'Constant pH Molecular Dynamics', 'Constant Redox Potential Molecular Dynamics', 'Continuous Constant pH Molecular Dynamics', 'NMR Refinement', 'X-ray and CryoEM Refinement', 'Locally Enhanced Sampling'])), allow_none=True) class EquilibrationSchema(Schema): @@ -159,13 +155,17 @@ class EquilibrationSchema(Schema): simulation_software = fields.List(fields.String(validate=OneOf(['Amber', 'GROMACS', 'LAMMPS', 'NAMD', 'OpenMM', 'CHARMM', 'DL_POLY', 'HOOMD-blue', 'Desmond', 'ACEMD', 'CP2K'])), allow_none=True) -class MinimisationSchema(Schema): - energy_tolerance = fields.Nested(ForceQuantitySchema, allow_none=True) - number_of_minimisation_steps = fields.Integer(allow_none=True) - minimisation_distance_step_size = fields.Nested(LengthQuantitySchema, allow_none=True) - minimisation_algorithm = fields.List(fields.String(validate=OneOf(['Steepest Descent', 'Conjugate Gradient', 'L-BFGS', 'XMIN', 'LMOD', 'None'])), allow_none=True) +class ProductionSchema(Schema): simulation_tool = fields.List(fields.String(validate=OneOf(['sander', 'pmemd', 'gem.pmemd', 'mdrun'])), allow_none=True) simulation_software = fields.List(fields.String(validate=OneOf(['Amber', 'GROMACS', 'LAMMPS', 'NAMD', 'OpenMM', 'CHARMM', 'DL_POLY', 'HOOMD-blue', 'Desmond', 'ACEMD', 'CP2K'])), allow_none=True) + simulation_software_version = SanitizedUnicode(allow_none=True) + simulation_method = fields.List(fields.String(validate=OneOf(['Self-guided Langevin Dynamics', 'Accelerated Molecular Dynamics', 'Gaussian Accelerated Molecular Dynamics', 'Targeted Molecular Dynamics', 'Nudged Elastic Band Calculations', 'Adaptive String Method', 'LMOD method', 'DL-FIND Optimization', 'Thermodynamic Integration (TI)', 'Linear Interaction Energies (LIE)', 'Replica Exchange Molecular Dynamics (REMD)', 'Adaptively Biased Molecular Dynamics (ABMD)', 'Steered Molecular Dynamics (SMD)', 'Umbrella Sampling', 'Metadynamics', 'Swarms of Trajectories String Method', 'Constant pH Molecular Dynamics', 'Constant Redox Potential Molecular Dynamics', 'Continuous Constant pH Molecular Dynamics', 'NMR Refinement', 'X-ray and CryoEM Refinement', 'Locally Enhanced Sampling'])), allow_none=True) + + +class AnalysisSchema(Schema): + analysis_tool = fields.List(fields.String(validate=OneOf(['mdout_analyzer.py', 'ambpdb', 'CPPTRAJ', 'PYTRAJ', 'MMPBSA.py', 'Free Energy Workflow (FEW)', 'edgember', 'SAX-RISM', 'SAX-MD', 'MoFT', 'ndfes', 'PLUMED', 'MDanalysis'])), allow_none=True) + analysis_software = fields.List(fields.String(validate=OneOf(['Visual Molecular Dynamics (VMD)', 'Schrödinger Maestro', 'PyMOL', 'Avogadro'])), allow_none=True) + analysis_method = fields.List(fields.String(validate=OneOf(['RMSD', 'DSSP', 'GIST', 'Hydrogen Bonds', 'Connolly surface', 'Radius of Gyration', 'BAR/PBSA'])), allow_none=True) class SimulationStagesSchema(Schema): @@ -176,6 +176,11 @@ class SimulationStagesSchema(Schema): analysis = fields.Nested(AnalysisSchema, allow_none=True) +class EnsembleSchema(Schema): + ensemble_type = fields.String(validate=OneOf(['NPT', 'NVT', 'NVE', 'μVT']), allow_none=True) + random_seed = fields.Integer(allow_none=True) + + class IntegratorSchema(Schema): integrator_algorithm = fields.String(validate=OneOf(['Velocity-Verlet', 'Leap-frog', 'Verlet', 'Euler']), allow_none=True) frame_step = fields.Nested(TimeQuantitySchema, allow_none=True) @@ -206,11 +211,6 @@ class ThermostatSchema(Schema): friction_coefficient = fields.Nested(FrictionCoefficientQuantitySchema, allow_none=True) -class EnsembleSchema(Schema): - ensemble_type = fields.String(validate=OneOf(['NPT', 'NVT', 'NVE', 'μVT']), allow_none=True) - random_seed = fields.Integer(allow_none=True) - - class InteractionsSchema(Schema): restraints = fields.Boolean(allow_none=True) electrostatic_cutoff_distance = fields.Nested(LengthQuantitySchema, allow_none=True) @@ -227,6 +227,13 @@ class SimulationSettingsSchema(Schema): interactions = fields.Nested(InteractionsSchema, allow_none=True) +class SystemCountsSchema(Schema): + total_molecule_count = fields.Integer(allow_none=True) + total_atom_count = fields.Integer(allow_none=True) + unique_molecule_count = fields.Integer(allow_none=True) + salt_concentration = fields.Nested(ConcentrationQuantitySchema, allow_none=True) + + class MoleculeIDSchema(Schema): PDB_ID = SanitizedUnicode(allow_none=True) UNIPROT_ID = SanitizedUnicode(allow_none=True) @@ -249,13 +256,6 @@ class MoleculeIDSchema(Schema): simulated_molecule_name = SanitizedUnicode(allow_none=True) -class SystemCountsSchema(Schema): - total_molecule_count = fields.Integer(allow_none=True) - total_atom_count = fields.Integer(allow_none=True) - unique_molecule_count = fields.Integer(allow_none=True) - salt_concentration = fields.Nested(ConcentrationQuantitySchema, allow_none=True) - - class SystemCompositionSchema(Schema): system_counts = fields.Nested(SystemCountsSchema, allow_none=True) molecule_ID = fields.List(fields.Nested(MoleculeIDSchema), allow_none=True) @@ -275,6 +275,11 @@ class SimulationObservablesSchema(Schema): simulation_averages = fields.Nested(SimulationAveragesSchema, allow_none=True) +class ConnectivitySchema(Schema): + bonds = fields.Boolean(allow_none=True) + dihedrals = fields.Boolean(allow_none=True) + + class ParticlesSchema(Schema): masses = fields.Boolean(allow_none=True) fixed_charges = fields.Boolean(allow_none=True) @@ -283,11 +288,6 @@ class ParticlesSchema(Schema): resolution = fields.String(validate=OneOf(['All Atom', 'United Atom', 'Coarse-Grained', 'Mesoscale']), allow_none=True) -class ConnectivitySchema(Schema): - bonds = fields.Boolean(allow_none=True) - dihedrals = fields.Boolean(allow_none=True) - - class TopologyMetadataSchema(Schema): connectivity = fields.Nested(ConnectivitySchema, allow_none=True) particles = fields.Nested(ParticlesSchema, allow_none=True) @@ -316,13 +316,13 @@ class TrajectoryMetadataSchema(Schema): trajectory_output = fields.Nested(TrajectoriesSchema, allow_none=True) -class NucleicPotentialSchema(Schema): - nucleic_potential_name = fields.String(validate=OneOf(['ff99-bsc0', 'ff99OL3', 'LJbb', 'ROC', 'Shaw', 'OL15', 'OL21', 'OL24', 'OL3', 'bsc1', 'terminal_monophosphate']), allow_none=True) +class WaterPotentialSchema(Schema): + water_potential_name = fields.String(validate=OneOf(['OPC', 'OPC3', 'OPC3POL', 'POL3', 'TIP3P', 'TIP3PFB', 'TIP4PFB', 'TIP4P', 'TIP5P', 'TIP4PEW', 'SPCE', 'SPCEB', 'SPC/Fw', 'q-SPC/Fw']), allow_none=True) modified = fields.Boolean(allow_none=True) -class WaterPotentialSchema(Schema): - water_potential_name = fields.String(validate=OneOf(['OPC', 'OPC3', 'OPC3POL', 'POL3', 'TIP3P', 'TIP3PFB', 'TIP4PFB', 'TIP4P', 'TIP5P', 'TIP4PEW', 'SPCE', 'SPCEB', 'SPC/Fw', 'q-SPC/Fw']), allow_none=True) +class ProteinPotentialSchema(Schema): + protein_potential_name = fields.String(validate=OneOf(['ff19SB', 'ff99SB', 'ff99SB-ILDN', 'ff99SB-disp', 'ff14SB', 'ff14SBonlysc', 'ff15ipq', 'fb15', 'ff03', 'ff03ua', 'phosaa10', 'phosaa14SB', 'phosaa19SB', 'ff14SB_modAA', 'ff19SB_modAA']), allow_none=True) modified = fields.Boolean(allow_none=True) @@ -331,8 +331,8 @@ class LipidPotentialSchema(Schema): modified = fields.Boolean(allow_none=True) -class ProteinPotentialSchema(Schema): - protein_potential_name = fields.String(validate=OneOf(['ff19SB', 'ff99SB', 'ff99SB-ILDN', 'ff99SB-disp', 'ff14SB', 'ff14SBonlysc', 'ff15ipq', 'fb15', 'ff03', 'ff03ua', 'phosaa10', 'phosaa14SB', 'phosaa19SB', 'ff14SB_modAA', 'ff19SB_modAA']), allow_none=True) +class NucleicPotentialSchema(Schema): + nucleic_potential_name = fields.String(validate=OneOf(['ff99-bsc0', 'ff99OL3', 'LJbb', 'ROC', 'Shaw', 'OL15', 'OL21', 'OL24', 'OL3', 'bsc1', 'terminal_monophosphate']), allow_none=True) modified = fields.Boolean(allow_none=True) @@ -367,18 +367,6 @@ class PotentialMetadataSchema(Schema): machine_learned_potential = fields.Nested(MachineLearnedPotentialSchema, allow_none=True) -class SoftwareSchema(Schema): - operating_system = fields.String(validate=OneOf(['Linux', 'macOS', 'Windows']), allow_none=True) - scheduler = fields.String(validate=OneOf(['SLURM', 'PBS', 'LSF', 'SGE', 'None']), allow_none=True) - MPI_library = fields.String(validate=OneOf(['OpenMPI', 'MPICH', 'IntelMPI', 'MVAPICH2', 'None']), allow_none=True) - container_runtime = fields.String(validate=OneOf(['Apptainer', 'Docker', 'Podman', 'None']), allow_none=True) - - -class PerformanceSchema(Schema): - wall_time = fields.Nested(TimeQuantitySchema, allow_none=True) - energy_consumption = fields.Nested(EnergyQuantitySchema, allow_none=True) - - class HardwareSchema(Schema): execution_platform = fields.String(validate=OneOf(['HPC Cluster', 'Cloud VM', 'Local']), allow_none=True) node_type = fields.String(validate=OneOf(['CPU only', 'GPU Accelerated', 'Hybrid CPU GPU']), allow_none=True) @@ -393,6 +381,18 @@ class HardwareSchema(Schema): memory_per_node = fields.Nested(ByteQuantitySchema, allow_none=True) +class SoftwareSchema(Schema): + operating_system = fields.String(validate=OneOf(['Linux', 'macOS', 'Windows']), allow_none=True) + scheduler = fields.String(validate=OneOf(['SLURM', 'PBS', 'LSF', 'SGE', 'None']), allow_none=True) + MPI_library = fields.String(validate=OneOf(['OpenMPI', 'MPICH', 'IntelMPI', 'MVAPICH2', 'None']), allow_none=True) + container_runtime = fields.String(validate=OneOf(['Apptainer', 'Docker', 'Podman', 'None']), allow_none=True) + + +class PerformanceSchema(Schema): + wall_time = fields.Nested(TimeQuantitySchema, allow_none=True) + energy_consumption = fields.Nested(EnergyQuantitySchema, allow_none=True) + + class ComputationalEnvironmentSchema(Schema): hardware = fields.Nested(HardwareSchema, allow_none=True) software = fields.Nested(SoftwareSchema, allow_none=True) From 0c47f28c356896a1c0f57f0e7b1efcd9891176f1 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Tue, 15 Sep 2026 17:24:49 +0100 Subject: [PATCH 3/5] look for nested lists and account for these --- .../fieldextraction/generate_marshmallow.py | 22 ++++++++++++++----- project/linkml/invenio/biosimdb_fields.py | 8 +++---- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/biosim_schema/fieldextraction/generate_marshmallow.py b/biosim_schema/fieldextraction/generate_marshmallow.py index 4298b8f..02f912e 100644 --- a/biosim_schema/fieldextraction/generate_marshmallow.py +++ b/biosim_schema/fieldextraction/generate_marshmallow.py @@ -83,9 +83,17 @@ def _field_and_mapping_for_slot(self, cls_name: str, slot_name: str): base_mapping = {"type": PRIMITIVE_MAPPING.get(range_name, "text")} depends = set() + # handle lists and nested lists if slot.multivalued: - base_field = f"fields.List({base_field})" - # ES: list of objects/scalars uses the same mapping as a single item + depth = 1 + if slot.array is not None: + depth = ( + slot.array.maximum_number_dimensions + or slot.array.minimum_number_dimensions + or 1 + ) + for _ in range(depth): + base_field = f"fields.List({base_field})" kwargs = [] if slot.required: @@ -191,9 +199,13 @@ def render_root(self, root_class: str | None = None) -> str: slot = self.view.induced_slot(slot_name, root) if self.view.get_class(slot.range) is not None: nested = schema_class_name(slot.range) - fields_lines.append( - f' "{slot_name}": fields.Nested({nested}, allow_none=True),' - ) + if slot.multivalued: + nested_field = ( + f"fields.List(fields.Nested({nested}), allow_none=True)" + ) + else: + nested_field = f"fields.Nested({nested}, allow_none=True)" + fields_lines.append(f' "{slot_name}": {nested_field},') else: _, field_expr, _, _, _, _ = self._field_and_mapping_for_slot( root, slot_name diff --git a/project/linkml/invenio/biosimdb_fields.py b/project/linkml/invenio/biosimdb_fields.py index 283d7eb..2a28134 100644 --- a/project/linkml/invenio/biosimdb_fields.py +++ b/project/linkml/invenio/biosimdb_fields.py @@ -123,17 +123,17 @@ class VectorTimeQuantitySchema(Schema): class MatrixPressureQuantitySchema(Schema): - vector_value = fields.List(fields.Float(), allow_none=True) + vector_value = fields.List(fields.List(fields.Float()), allow_none=True) value_unit = fields.String(validate=OneOf(['bar', 'Pa']), allow_none=True) class MatrixCompressibilityQuantitySchema(Schema): - vector_value = fields.List(fields.Float(), allow_none=True) + vector_value = fields.List(fields.List(fields.Float()), allow_none=True) value_unit = fields.String(validate=OneOf(['1/bar', '1/Pa']), allow_none=True) class MatrixQuantitySchema(Schema): - vector_value = fields.List(fields.Float(), allow_none=True) + vector_value = fields.List(fields.List(fields.Float()), allow_none=True) value_unit = SanitizedUnicode(allow_none=True) @@ -429,7 +429,7 @@ class SimulationMetadataSchema(Schema): "composition": fields.Nested(SystemCompositionSchema, allow_none=True), "potentials": fields.Nested(PotentialMetadataSchema, allow_none=True), "compute": fields.Nested(ComputationalEnvironmentSchema, allow_none=True), - "files": fields.Nested(FileMetadataSchema, allow_none=True), + "files": fields.List(fields.Nested(FileMetadataSchema), allow_none=True), "biosim_schema_version": SanitizedUnicode(allow_none=True), } From f1739a5c295dabe8d240b7d8b8b5e3db102addae Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Wed, 16 Sep 2026 15:18:24 +0100 Subject: [PATCH 4/5] update tests to include generate_invenio_marshmallow as expected generator --- tests/test_utils/test_generate_artifacts.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_utils/test_generate_artifacts.py b/tests/test_utils/test_generate_artifacts.py index aa9f3b1..1047f63 100644 --- a/tests/test_utils/test_generate_artifacts.py +++ b/tests/test_utils/test_generate_artifacts.py @@ -132,10 +132,13 @@ def test_generate_derived_calls_expected_generators(monkeypatch, tmp_path): mod, "generate_engine_mappings", lambda p: calls.append("mappings") ) monkeypatch.setattr(mod, "generate_summary", lambda p: calls.append("summary")) + monkeypatch.setattr( + mod, "generate_invenio_marshmallow", lambda p: calls.append("marshmallow") + ) mod.generate_derived(paths) - assert calls == ["webform", "mappings", "summary"] + assert calls == ["webform", "mappings", "summary", "marshmallow"] def test_main_dispatches_jsonld(monkeypatch, tmp_path): From f20eb9f8eb51e835f94cec0ba470bba0743078c4 Mon Sep 17 00:00:00 2001 From: Jas Kalayan Date: Wed, 16 Sep 2026 20:24:39 +0100 Subject: [PATCH 5/5] add test for gen marshmallow schema --- .../test_generate_marshmallow.py | 153 ++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 tests/test_fieldextraction/test_generate_marshmallow.py diff --git a/tests/test_fieldextraction/test_generate_marshmallow.py b/tests/test_fieldextraction/test_generate_marshmallow.py new file mode 100644 index 0000000..b0c5fdc --- /dev/null +++ b/tests/test_fieldextraction/test_generate_marshmallow.py @@ -0,0 +1,153 @@ +import sys +from types import SimpleNamespace + +import pytest + +from biosim_schema.fieldextraction import generate_marshmallow as mod + + +class View: + def __init__(self): + self.slots = { + ("Root", "child"): SimpleNamespace( + range="Child", multivalued=False, array=None, required=True + ), + ("Root", "tags"): SimpleNamespace( + range="string", + multivalued=True, + array=SimpleNamespace(maximum_number_dimensions=2), + required=False, + ), + ("Child", "kind"): SimpleNamespace( + range="Kind", multivalued=False, array=None, required=False + ), + ("Child", "count"): SimpleNamespace( + range="integer", multivalued=False, array=None, required=True + ), + } + self.classes = { + "Root": SimpleNamespace(tree_root=True), + "Child": SimpleNamespace(tree_root=False), + } + self.enums = { + "Kind": SimpleNamespace(permissible_values={"a": None, "b": None}) + } + + def induced_slot(self, name, cls): + return self.slots[(cls, name)] + + def get_class(self, name): + return self.classes.get(name) + + def get_enum(self, name): + return self.enums.get(name) + + def class_slots(self, name): + return [key[1] for key in self.slots if key[0] == name] + + def all_classes(self, imports=True): + return self.classes + + +def generator(): + result = object.__new__(mod.MarshmallowGenerator) + result.view = View() + result.builds = {} + return result + + +def test_field_rendering_covers_nested_enum_and_lists(): + """Renders nested, enum, primitive, required, and list fields.""" + gen = generator() + + nested = gen._field_and_mapping_for_slot("Root", "child") + listed = gen._field_and_mapping_for_slot("Root", "tags") + enum = gen._field_and_mapping_for_slot("Child", "kind") + + assert "Nested(ChildSchema" in nested[1] + assert nested[3] == {"Child"} + assert listed[1].count("fields.List") == 2 + assert "allow_none=True" in listed[1] + assert "OneOf" in enum[1] + assert enum[2] == {"type": "keyword"} + + +def test_generate_renders_nested_schema_and_mapping(): + """Generates schemas, mappings, and nested root fields.""" + gen = generator() + + code = gen.generate() + + assert "class ChildSchema(Schema)" in code + assert "class RootSchema(Schema)" in code + assert '"child": fields.Nested(ChildSchema, allow_none=True)' in code + assert '"tags": fields.List(fields.List(SanitizedUnicode()' in code + assert "FIELDS =" in code + assert "MAPPING =" in code + + +def test_recursive_nested_mapping_breaks_cycle(): + """Stops recursive nested mappings.""" + gen = generator() + gen.builds = { + "Root": mod.ClassBuild( + "Root", + mapping={"child": {"type": "object", "_nested_class": "Child"}}, + depends_on=["Child"], + ), + "Child": mod.ClassBuild( + "Child", + mapping={"root": {"type": "object", "_nested_class": "Root"}}, + depends_on=["Root"], + ), + } + + gen._resolve_nested_mappings() + + assert gen.builds["Child"].mapping["root"] == {"type": "object"} + + +def test_ordering_and_root_errors(): + """Orders dependencies and rejects schemas without a root.""" + gen = generator() + gen.builds = { + "Root": mod.ClassBuild("Root", depends_on=["Child"]), + "Child": mod.ClassBuild("Child"), + } + + assert gen._topo_order() == ["Child", "Root"] + + gen.view.classes = {"Other": SimpleNamespace(tree_root=False)} + with pytest.raises(ValueError, match="tree_root"): + gen._find_tree_root() + + +def test_main_writes_generated_code(monkeypatch, tmp_path): + """CLI writes generated output to the requested file.""" + output = tmp_path / "generated.py" + + class FakeGenerator: + def __init__(self, schema): + assert schema == "schema.yaml" + + def generate(self, root_class=None): + assert root_class == "Root" + return "generated" + + monkeypatch.setattr(mod, "MarshmallowGenerator", FakeGenerator) + monkeypatch.setattr( + sys, + "argv", + [ + "generate_marshmallow", + "schema.yaml", + "-o", + str(output), + "--root-class", + "Root", + ], + ) + + mod.main() + + assert output.read_text() == "generated"