diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py index 941efd61c2..0cb555e4d1 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: 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) 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 = [] 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)