argcheck: restrict the type of elements in a list (#179)

Signed-off-by: Jinzhe Zeng <[email protected]>

dpgen2/entrypoint/args.py

@@ -1,4 +1,7 @@
 import textwrap
+from typing import (
+    List,
+)
 
 import dargs
 from dargs import (
@@ -53,7 +56,7 @@ def dp_dist_train_args():
             doc=doc_config,
         ),
         Argument(
-            "template_script", [list, str], optional=False, doc=doc_template_script
+            "template_script", [List[str], str], optional=False, doc=doc_template_script
         ),
         Argument("student_model_path", str, optional=True, doc=dock_student_model_path),
     ]
@@ -76,11 +79,11 @@ def dp_train_args():
         ),
         Argument("numb_models", int, optional=True, default=4, doc=doc_numb_models),
         Argument(
-            "template_script", [list, str], optional=False, doc=doc_template_script
+            "template_script", [List[str], str], optional=False, doc=doc_template_script
         ),
         Argument(
             "init_models_paths",
-            list,
+            List[str],
             optional=True,
             default=None,
             doc=doc_init_models_paths,
@@ -175,7 +178,7 @@ def lmp_args():
         ),
         Argument(
             "convergence",
-            list,
+            dict,
             [],
             [variant_conv()],
             optional=False,
@@ -191,7 +194,7 @@ def lmp_args():
             doc=doc_configuration,
             alias=["configuration"],
         ),
-        Argument("stages", list, optional=False, doc=doc_stages),
+        Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
     ]
 
 
@@ -266,8 +269,8 @@ def input_args():
     doc_init_sys = "The inital data systems"
 
     return [
-        Argument("type_map", list, optional=False, doc=doc_type_map),
-        Argument("mass_map", list, optional=False, doc=doc_mass_map),
+        Argument("type_map", List[str], optional=False, doc=doc_type_map),
+        Argument("mass_map", List[float], optional=False, doc=doc_mass_map),
         Argument(
             "init_data_prefix",
             str,
@@ -280,7 +283,11 @@ def input_args():
             "do_finetune", bool, optional=True, default=False, doc=doc_do_finetune
         ),
         Argument(
-            "init_data_sys", [list, str], optional=False, default=None, doc=doc_init_sys
+            "init_data_sys",
+            [List[str], str],
+            optional=False,
+            default=None,
+            doc=doc_init_sys,
         ),
     ]
 
@@ -479,7 +486,7 @@ def submit_args(default_step_config=normalize_step_dict({})):
             ),
             Argument(
                 "upload_python_packages",
-                [list, str],
+                [List[str], str],
                 optional=True,
                 default=None,
                 doc=doc_upload_python_packages,

examples/almg/input.json

@@ -2,7 +2,8 @@
     "bohrium_config": {
 	"username":	"__write your bohrium username",
 	"password":	"__write your bohrium password",
-	"project_id" :	"__write your bohrium project ID",
+	"_project_id" :	"__write your bohrium project ID",
+	"project_id": 123456,
 	"_comment" : "all"
     },
 

examples/chno/input.json

@@ -2,7 +2,8 @@
     "bohrium_config": {
 	"username":	"__write your bohrium username",
 	"password":	"__write your bohrium password",
-	"project_id" :	"__write your bohrium project ID",
+	"_project_id" :	"__write your bohrium project ID",
+	"project_id" :	123456,
 	"_comment" : "all"
     },
 

pyproject.toml

@@ -49,7 +49,7 @@ docs = [
      'myst_parser',
      'deepmodeling_sphinx',
      'sphinx-argparse',
-     "dargs>=0.3.1",
+     "dargs>=0.4.1",
 ]
 test = [
      'fakegaussian>=0.0.3',

tests/check_examples.py

@@ -0,0 +1,29 @@
+import json
+import unittest
+from pathlib import (
+    Path,
+)
+
+from dpgen2.entrypoint.args import (
+    normalize,
+)
+
+p_examples = Path(__file__).parent.parent / "examples"
+
+input_files = (
+    p_examples / "almg" / "input.json",
+    # p_examples / "almg" / "input-v005.json",
+    # p_examples / "almg" / "dp_template.json",
+    p_examples / "ch4" / "input_dist.json",
+    # p_examples / "chno" / "dpa_manyi.json",
+    p_examples / "chno" / "input.json",
+)
+
+
+class TestExamples(unittest.TestCase):
+    def test_arguments(self):
+        for fn in input_files:
+            with self.subTest(fn=fn):
+                with open(fn) as f:
+                    jdata = json.load(f)
+                normalize(jdata)

tests/entrypoint/test_submit_args.py

@@ -51,7 +51,7 @@ def test(self):
             new_data["bohrium_config"],
             None,
         )
-        self.assertEqual(old_data["model_devi_jobs"], new_data["explore"]["stages"])
+        self.assertEqual(old_data["model_devi_jobs"], new_data["explore"]["stages"][0])
         new_data["explore"]["configurations"][0].pop("type")
         self.assertEqual(old_data["sys_configs"], new_data["explore"]["configurations"])
         self.assertEqual(
@@ -413,10 +413,10 @@ def test_bohrium(self):
 		"concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]
 	    }
 	],
-	"stages":	[
+	"stages":	[[
 	    { "_idx": 0, "ensemble": "nvt", "nsteps": 20, "press": [1.0,2.0], "sys_idx": [0], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 
 	    }
-	],
+	]],
 	"_comment" : "all"
     },
     "fp" : {