feat: support opt-in $ref loading in workflows

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.

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.

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
 {

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"

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")
 
 

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

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")

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
 
 

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',

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)