from argparse import ArgumentParser from pathlib import Path from linkml_runtime.dumpers import yaml_dumper from nwb_linkml.generators.pydantic import NWBPydanticGenerator from nwb_linkml.src.nwb_linkml import io def generate_core_yaml(output_path:Path): core = io.load_nwb_core() built_schemas = core.build().schemas for schema in built_schemas: output_file = output_path / (schema.name + '.yaml') yaml_dumper.dump(schema, output_file) def generate_core_pydantic(yaml_path:Path, output_path:Path): for schema in yaml_path.glob('*.yaml'): python_name = schema.stem.replace('.', '_').replace('-', '_') pydantic_file = (output_path / python_name).with_suffix('.py') generator = NWBPydanticGenerator( str(schema), pydantic_version='2', emit_metadata=True, gen_classvars=True, gen_slots=True ) gen_pydantic = generator.serialize() with open(pydantic_file, 'w') as pfile: pfile.write(gen_pydantic) def parser() -> ArgumentParser: parser = ArgumentParser('Generate NWB core schema') parser.add_argument( '--yaml', help="directory to export linkML schema to", type=Path, default=Path(__file__).parent.parent / 'nwb_linkml' / 'src' / 'nwb_linkml' / 'schema' ) parser.add_argument( '--pydantic', help="directory to export pydantic models", type=Path, default=Path(__file__).parent.parent / 'nwb_linkml' / 'src' / 'nwb_linkml' / 'models' ) return parser def main(): args = parser().parse_args() args.yaml.mkdir(exist_ok=True) args.pydantic.mkdir(exist_ok=True) generate_core_yaml(args.yaml) generate_core_pydantic(args.yaml, args.pydantic) if __name__ == "__main__": main()