From 337fd45cea9b760c2349982cacfa6269156676d2 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Mon, 26 Jul 2021 04:54:51 -0400
Subject: [PATCH 1/3] enhance the cli to generate doc json file

---
 deepmd/entrypoints/doc.py  | 11 ++++++++---
 deepmd/entrypoints/main.py | 10 ++++++++--
 deepmd/utils/argcheck.py   | 10 +++++++++-
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py
index 941efd61c2..376a08ce8e 100644
--- a/deepmd/entrypoints/doc.py
+++ b/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: bool = "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)
diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py
index cda9856a24..9557976bf6 100644
--- a/deepmd/entrypoints/main.py
+++ b/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":
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 4dfd552d89..3f75b39394 100644
--- a/deepmd/utils/argcheck.py
+++ b/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 = []

From 3a5afd32ad181d7ad78677394b0e2ba30a8dea0a Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 27 Jul 2021 03:26:21 -0400
Subject: [PATCH 2/3] bump dargs version; add argument to tests

---
 requirements.txt                     | 2 +-
 source/tests/test_argument_parser.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index e3a8f501ab..21befa3722 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 numpy
 scipy
 pyyaml
-dargs >= 0.2.2
+dargs >= 0.2.6
 typing_extensions; python_version < "3.7"
diff --git a/source/tests/test_argument_parser.py b/source/tests/test_argument_parser.py
index f9f28fb81b..615199603a 100644
--- a/source/tests/test_argument_parser.py
+++ b/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)
     

From 4770db57bb973c1103172b633f3589486cef4cb6 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sat, 31 Jul 2021 04:04:26 -0400
Subject: [PATCH 3/3] correct the type hint of `out_type`

---
 deepmd/entrypoints/doc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py
index 376a08ce8e..0cb555e4d1 100644
--- a/deepmd/entrypoints/doc.py
+++ b/deepmd/entrypoints/doc.py
@@ -5,7 +5,7 @@
 __all__ = ["doc_train_input"]
 
 
-def doc_train_input(*, out_type: bool = "rst", **kwargs):
+def doc_train_input(*, out_type: str = "rst", **kwargs):
     """Print out trining input arguments to console."""
     if out_type == "rst":
         doc_str = gen_doc(make_anchor=True)