feat: support generating JSON schema for integration with VSCode

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

deepmd/entrypoints/doc.py

@@ -4,6 +4,7 @@
 from deepmd.utils.argcheck import (
     gen_doc,
     gen_json,
+    gen_json_schema,
 )
 
 __all__ = ["doc_train_input"]
@@ -15,6 +16,8 @@ def doc_train_input(*, out_type: str = "rst", **kwargs):
         doc_str = gen_doc(make_anchor=True)
     elif out_type == "json":
         doc_str = gen_json()
+    elif out_type == "json_schema":
+        doc_str = gen_json_schema()
     else:
         raise RuntimeError(f"Unsupported out type {out_type}")
     print(doc_str)  # noqa: T201

deepmd/main.py

@@ -500,7 +500,7 @@ def main_parser() -> argparse.ArgumentParser:
     parsers_doc.add_argument(
         "--out-type",
         default="rst",
-        choices=["rst", "json"],
+        choices=["rst", "json", "json_schema"],
         type=str,
         help="The output type",
     )

deepmd/utils/argcheck.py

@@ -14,7 +14,13 @@
     Variant,
     dargs,
 )
+from dargs.json_schema import (
+    generate_json_schema,
+)
 
+from deepmd import (
+    __version__,
+)
 from deepmd.common import (
     VALID_ACTIVATION,
     VALID_PRECISION,
@@ -2450,6 +2456,18 @@ def gen_args(**kwargs) -> List[Argument]:
     ]
 
 
+def gen_json_schema() -> str:
+    """Generate JSON schema.
+
+    Returns
+    -------
+    str
+        JSON schema.
+    """
+    arg = Argument("DeePMD-kit", dict, gen_args(), doc=f"DeePMD-kit {__version__}")
+    return json.dumps(generate_json_schema(arg))
+
+
 def normalize(data):
     base = Argument("base", dict, gen_args())
     data = base.normalize_value(data, trim_pattern="_*")

doc/train/train-input.rst

@@ -1,8 +1,42 @@
 Training Parameters
 ======================================
 .. note::
-   One can load, modify, and export the input file by using our effective web-based tool `DP-GUI <https://deepmodeling.com/dpgui/input/deepmd-kit-2.0>`_ online or hosted using the :ref:`command line interface <cli>` :code:`dp gui`. All training parameters below can be set in DP-GUI. By clicking "SAVE JSON", one can download the input file for furthur training.
+   One can load, modify, and export the input file by using our effective web-based tool `DP-GUI <https://deepmodeling.com/dpgui/input/deepmd-kit-2.0>`_ online or hosted using the :ref:`command line interface <cli>` :code:`dp gui`. All training parameters below can be set in DP-GUI. By clicking "SAVE JSON", one can download the input file for further training.
+
+.. note::
+   One can benefit from IntelliSense and validation when
+   :ref:`writing JSON files using Visual Studio Code <json_vscode>`.
+   See :ref:`here <json_vscode>` to learn how to configure.
 
 .. dargs::
-   :module: deepmd.tf.utils.argcheck
+   :module: deepmd.utils.argcheck
    :func: gen_args
+
+.. _json_vscode:
+
+Writing JSON files using Visual Studio Code
+-------------------------------------------
+
+When writing JSON files using Visual Studio Code, one can benefit from IntelliSense and
+validation by adding a `JSON schema <https://json-schema.org/>`_.
+To do so, in a VS Code workspace, one can generate a JSON schema file for the input file by running the following command:
+
+.. code-block:: bash
+
+   dp doc-train-input --out-type json_schema > deepmd.json
+
+Then one can `map the schema <https://code.visualstudio.com/docs/languages/json#_mapping-to-a-schema-in-the-workspace>`_
+by updating the workspace settings in the `.vscode/settings.json` file as follows:
+
+.. code-block:: json
+
+   {
+      "json.schemas": [
+         {
+               "fileMatch": [
+                  "/**/*.json"
+               ],
+               "url": "./deepmd.json"
+         }
+      ]
+   }

pyproject.toml

@@ -39,7 +39,7 @@ dependencies = [
     'numpy',
     'scipy',
     'pyyaml',
-    'dargs >= 0.4.1',
+    'dargs >= 0.4.6',
     'typing_extensions; python_version < "3.8"',
     'importlib_metadata>=1.4; python_version < "3.8"',
     'h5py',

source/tests/common/test_doc_train_input.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import io
+import json
+import unittest
+from contextlib import (
+    redirect_stdout,
+)
+
+from deepmd.entrypoints.doc import (
+    doc_train_input,
+)
+
+
+class TestDocTrainInput(unittest.TestCase):
+    def test_rst(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            doc_train_input(out_type="rst")
+        self.assertTrue(f.getvalue() != "")
+
+    def test_json(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            doc_train_input(out_type="json")
+        # validate json
+        json.loads(f.getvalue())
+
+    def test_json_schema(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            doc_train_input(out_type="json_schema")
+        # validate json
+        json.loads(f.getvalue())