diff --git a/doc/init/init-reaction.md b/doc/init/init-reaction.md index c7366dcbf..c6627b2e9 100644 --- a/doc/init/init-reaction.md +++ b/doc/init/init-reaction.md @@ -18,4 +18,10 @@ An example of `reaction.json` is given below: For detailed parameters, see [parametes](init-reaction-jdata.rst) and [machine parameters](init-reaction-mdata.rst). +To enable loading external JSON/YAML snippets via `$ref` in the reaction parameter file, add `--allow-ref` explicitly (disabled by default for security): + +```bash +dpgen init_reaction reaction.json machine.json --allow-ref +``` + The genereated data can be used to continue DP-GEN concurrent learning workflow. Read [Energy & Fuels, 2021, 35 (1), 762–769](https://10.1021/acs.energyfuels.0c03211) for details. diff --git a/doc/run/overview-of-the-run-process.md b/doc/run/overview-of-the-run-process.md index 590312ff3..7fb63b356 100644 --- a/doc/run/overview-of-the-run-process.md +++ b/doc/run/overview-of-the-run-process.md @@ -16,6 +16,12 @@ Here, we give a general description of the run process. We can execute the run p dpgen run param.json machine.json ``` +To enable loading external JSON/YAML snippets via `$ref` in `param.json`, add `--allow-ref` explicitly (disabled by default for security): + +```sh +dpgen run param.json machine.json --allow-ref +``` + The following files or folders will be created and upgraded by codes： - iter.00000x contains the main results that DP-GEN generates in the first iteration. diff --git a/doc/simplify/simplify.md b/doc/simplify/simplify.md index 4dd07a3eb..d8798ea57 100644 --- a/doc/simplify/simplify.md +++ b/doc/simplify/simplify.md @@ -6,6 +6,12 @@ Use the following script to start the workflow: dpgen simplify param.json machine.json ``` +To enable loading external JSON/YAML snippets via `$ref`, add `--allow-ref` explicitly (disabled by default for security): + +```bash +dpgen simplify param.json machine.json --allow-ref +``` + Here is an example of `param.json` for QM7 dataset: ```json { diff --git a/dpgen/data/reaction.py b/dpgen/data/reaction.py index 46cde2ff8..754f1a425 100644 --- a/dpgen/data/reaction.py +++ b/dpgen/data/reaction.py @@ -212,7 +212,7 @@ def gen_init_reaction(args): mdata = load_file(args.MACHINE) jdata_arginfo = init_reaction_jdata_arginfo() - jdata = normalize(jdata_arginfo, jdata) + jdata = normalize(jdata_arginfo, jdata, allow_ref=args.allow_ref) mdata = convert_mdata(mdata, ["reaxff", "build", "fp"]) record = "record.reaction" diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 46ec6e1ce..1dde3841f 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -4972,12 +4972,12 @@ def set_version(mdata): return mdata -def run_iter(param_file, machine_file): +def run_iter(param_file, machine_file, allow_ref: bool = False): jdata = load_file(param_file) mdata = load_file(machine_file) jdata_arginfo = run_jdata_arginfo() - jdata = normalize(jdata_arginfo, jdata, strict_check=False) + jdata = normalize(jdata_arginfo, jdata, strict_check=False, allow_ref=allow_ref) update_mass_map(jdata) @@ -5718,7 +5718,7 @@ def gen_run(args): if args.debug: dlog.setLevel(logging.DEBUG) dlog.info("start running") - run_iter(args.PARAM, args.MACHINE) + run_iter(args.PARAM, args.MACHINE, allow_ref=args.allow_ref) dlog.info("finished") diff --git a/dpgen/main.py b/dpgen/main.py index a0ac004e5..ecef58bd5 100644 --- a/dpgen/main.py +++ b/dpgen/main.py @@ -101,6 +101,11 @@ def main_parser() -> argparse.ArgumentParser: nargs="?", help="machine file, json/yaml format", ) + parser_init_reaction.add_argument( + "--allow-ref", + action="store_true", + help="Allow loading external JSON/YAML snippets through `$ref`. Disabled by default for security.", + ) parser_init_reaction.set_defaults(func=gen_init_reaction) # run @@ -110,6 +115,11 @@ def main_parser() -> argparse.ArgumentParser: parser_run.add_argument("PARAM", type=str, help="parameter file, json/yaml format") parser_run.add_argument("MACHINE", type=str, help="machine file, json/yaml format") parser_run.add_argument("-d", "--debug", action="store_true", help="log debug info") + parser_run.add_argument( + "--allow-ref", + action="store_true", + help="Allow loading external JSON/YAML snippets through `$ref`. Disabled by default for security.", + ) parser_run.set_defaults(func=gen_run) # run/report @@ -181,6 +191,11 @@ def main_parser() -> argparse.ArgumentParser: parser_run.add_argument("PARAM", type=str, help="parameter file, json/yaml format") parser_run.add_argument("MACHINE", type=str, help="machine file, json/yaml format") parser_run.add_argument("-d", "--debug", action="store_true", help="log debug info") + parser_run.add_argument( + "--allow-ref", + action="store_true", + help="Allow loading external JSON/YAML snippets through `$ref`. Disabled by default for security.", + ) parser_run.set_defaults(func=gen_simplify) # test diff --git a/dpgen/simplify/simplify.py b/dpgen/simplify/simplify.py index 2bc88e1f8..ec56f28cb 100644 --- a/dpgen/simplify/simplify.py +++ b/dpgen/simplify/simplify.py @@ -508,7 +508,7 @@ def make_fp(iter_index, jdata, mdata): make_fp_calculation(iter_index, jdata, mdata) -def run_iter(param_file, machine_file): +def run_iter(param_file, machine_file, allow_ref: bool = False): """Init (iter 0): init_pick. tasks (iter > 0): @@ -521,12 +521,22 @@ def run_iter(param_file, machine_file): 06 make_fp 07 run_fp (same as generator) 08 post_fp (same as generator) + + Parameters + ---------- + param_file : str + parameter file (json/yaml) + machine_file : str + machine file (json/yaml) + allow_ref : bool, default=False + Whether to allow loading external JSON/YAML snippets via ``$ref``. + Disabled by default for security. """ jdata = load_file(param_file) mdata = load_file(machine_file) jdata_arginfo = simplify_jdata_arginfo() - jdata = normalize(jdata_arginfo, jdata) + jdata = normalize(jdata_arginfo, jdata, allow_ref=allow_ref) # set up electron temperature use_ele_temp = jdata.get("use_ele_temp", 0) @@ -626,5 +636,5 @@ def gen_simplify(args): if args.debug: dlog.setLevel(logging.DEBUG) dlog.info("start simplifying") - run_iter(args.PARAM, args.MACHINE) + run_iter(args.PARAM, args.MACHINE, allow_ref=args.allow_ref) dlog.info("finished") diff --git a/dpgen/util.py b/dpgen/util.py index 73453d74b..c222a5539 100644 --- a/dpgen/util.py +++ b/dpgen/util.py @@ -78,7 +78,12 @@ def expand_sys_str(root_dir: Union[str, Path]) -> list[str]: return matches -def normalize(arginfo: Argument, data: dict, strict_check: bool = True) -> dict: +def normalize( + arginfo: Argument, + data: dict, + strict_check: bool = True, + allow_ref: bool = False, +) -> dict: """Normalize and check input data. Parameters @@ -89,14 +94,17 @@ def normalize(arginfo: Argument, data: dict, strict_check: bool = True) -> dict: input data strict_check : bool, default=True strict check data or not + allow_ref : bool, default=False + Whether to allow loading external JSON/YAML snippets via ``$ref``. + Disabled by default for security. Returns ------- dict normalized data """ - data = arginfo.normalize_value(data, trim_pattern="_*") - arginfo.check_value(data, strict=strict_check) + data = arginfo.normalize_value(data, trim_pattern="_*", allow_ref=allow_ref) + arginfo.check_value(data, strict=strict_check, allow_ref=allow_ref) return data diff --git a/pyproject.toml b/pyproject.toml index 36c4575e3..b3ae59d80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ 'GromacsWrapper>=0.9.0; python_version >= "3.12"', 'dpdispatcher>=0.3.11', 'netCDF4', - 'dargs>=0.4.0', + 'dargs>=0.5.0', 'h5py', 'pymatgen-analysis-defects', 'openbabel-wheel', diff --git a/tests/test_cli.py b/tests/test_cli.py index 0c2125d2a..da7709e9c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,10 +1,11 @@ import subprocess as sp +import sys import unittest class TestCLI(unittest.TestCase): def test_cli(self): - sp.check_output(["dpgen", "-h"]) + sp.check_output([sys.executable, "-m", "dpgen.main", "-h"]) for subcommand in ( "run", "simplify", @@ -13,4 +14,8 @@ def test_cli(self): "init_reaction", "autotest", ): - sp.check_output(["dpgen", subcommand, "-h"]) + output = sp.check_output( + [sys.executable, "-m", "dpgen.main", subcommand, "-h"] + ) + if subcommand in ("run", "simplify", "init_reaction"): + self.assertIn(b"--allow-ref", output)