enhance the cli to generate doc json file

deepmd/entrypoints/doc.py

@@ -1,11 +1,16 @@
 """Module that prints train input arguments docstrings."""
 
-from deepmd.utils.argcheck import gen_doc
+from deepmd.utils.argcheck import gen_doc, gen_json
 
 __all__ = ["doc_train_input"]
 
 
-def doc_train_input():
+def doc_train_input(*, out_type: str = "rst", **kwargs):
     """Print out trining input arguments to console."""
-    doc_str = gen_doc(make_anchor=True)
+    if out_type == "rst":
+        doc_str = gen_doc(make_anchor=True)
+    elif out_type == "json":
+        doc_str = gen_json()
+    else:
+        raise RuntimeError("Unsupported out type %s" % out_type)
     print(doc_str)

deepmd/entrypoints/main.py

@@ -313,12 +313,18 @@ def parse_args(args: Optional[List[str]] = None):
     )
 
     # * print docs script **************************************************************
-    subparsers.add_parser(
+    parsers_doc = subparsers.add_parser(
         "doc-train-input",
         parents=[parser_log],
         help="print the documentation (in rst format) of input training parameters.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
+    parsers_doc.add_argument(
+        "--out-type", 
+        default="rst", 
+        type=str, 
+        help="The output type"
+    )
 
     # * make model deviation ***********************************************************
     parser_model_devi = subparsers.add_parser(
@@ -428,7 +434,7 @@ def main():
     elif args.command == "compress":
         compress(**dict_args)
     elif args.command == "doc-train-input":
-        doc_train_input()
+        doc_train_input(**dict_args)
     elif args.command == "model-devi":
         make_model_devi(**dict_args)
     elif args.command == "convert-from":

deepmd/utils/argcheck.py

@@ -1,5 +1,6 @@
-from dargs import dargs, Argument, Variant
+from dargs import dargs, Argument, Variant, ArgumentEncoder
 from deepmd.common import ACTIVATION_FN_DICT, PRECISION_DICT
+import json
 
 
 def list_to_doc(xx):
@@ -622,6 +623,13 @@ def gen_doc(*, make_anchor=True, make_link=True, **kwargs):
 
     return "\n\n".join(ptr)
 
+def gen_json(**kwargs):
+    return json.dumps((
+        model_args(),
+        learning_rate_args(),
+        loss_args(),
+        training_args(),
+    ), cls=ArgumentEncoder)
 
 def normalize_hybrid_list(hy_list):
     new_list = []

requirements.txt

@@ -1,5 +1,5 @@
 numpy
 scipy
 pyyaml
-dargs >= 0.2.2
+dargs >= 0.2.6
 typing_extensions; python_version < "3.7"

source/tests/test_argument_parser.py

@@ -282,7 +282,9 @@ def test_parser_compress(self):
 
     def test_parser_doc(self):
         """Test doc subparser."""
-        ARGS = {}
+        ARGS = {
+            "--out-type": dict(type=str, value="rst"),
+        }
 
         self.run_test(command="doc-train-input", mapping=ARGS)