CI test: check if documentation is complete (#1070)

Summary:
Fixes #1047

Basic but should do the trick. Not sure about the Github/actions part (never used this before)

### Changes

- Add CI test that checks if all exported dps have documentation

Pull Request resolved: #1070

Reviewed By: NivekT

Differential Revision: D43876159

Pulled By: ejguan

fbshipit-source-id: 3b77416ca8f082cff90abec981aa00ddc4232db8

Author: SvenDS9

.github/scripts/check_complete_doc.py

@@ -0,0 +1,67 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import sys
+
+
+def collect_init_dps(init_file_location):
+    init_dps = set()
+    with open(init_file_location) as init_file:
+        while (line := init_file.readline()) != "":
+            if line.startswith("__all__ "):
+                while (line := init_file.readline()) != "" and (stripped_line := line.strip()).startswith('"'):
+                    init_dps.add(stripped_line.replace(",", "").replace('"', ""))
+                break
+    return init_dps
+
+
+def collect_rst_dps(rst_file_location):
+    rst_dps = set()
+    with open(rst_file_location) as rst_file:
+        while (line := rst_file.readline()) != "":
+            if line.count("class_template.rst") > 0 or line.count("function.rst") > 0:
+                rst_file.readline()
+                while (line := rst_file.readline()) != "" and len(stripped_line := line.strip()) > 1:
+                    rst_dps.add(stripped_line)
+    return rst_dps
+
+
+def compare_sets(set_a, set_b, ignore_set=None):
+    res = set_a.difference(set_b)
+    if ignore_set is not None:
+        res.difference_update(ignore_set)
+    return res
+
+
+def main():
+    datapipes_folder = os.path.join("torchdata", "datapipes")
+    init_file = "__init__.py"
+    docs_source_folder = os.path.join("docs", "source")
+    exit_code = 0
+
+    for target, ignore_set in zip(["iter", "map", "utils"], [{"IterDataPipe", "Extractor"}, {"MapDataPipe"}, {}]):
+        init_path = os.path.join(datapipes_folder, target, init_file)
+        rst_path = os.path.join(docs_source_folder, "torchdata.datapipes." + target + ".rst")
+
+        init_set = collect_init_dps(init_path)
+        rst_set = collect_rst_dps(rst_path)
+
+        dif_init = compare_sets(init_set, rst_set, ignore_set)
+        dif_rst = compare_sets(rst_set, init_set)
+
+        for elem in dif_init:
+            print(f"Please add {elem} to {rst_path}")
+            exit_code = 1
+        for elem in dif_rst:
+            print(f"{elem} is present in {rst_path} but not in {init_path}")
+            exit_code = 1
+
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/lint.yml

@@ -73,3 +73,16 @@ jobs:
             echo 'Please fix the above mypy warnings.'
             false
           fi
+
+  complete_documentation:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    runs-on: ubuntu-latest
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.8"
+      - name: Check out source repository
+        uses: actions/checkout@v3
+      - name: Check if documentation is complete
+        run: python ./.github/scripts/check_complete_doc.py

docs/source/torchdata.datapipes.utils.rst

@@ -24,6 +24,7 @@ Common Utility Functions
     :toctree: generated/
     :template: function.rst
 
+    janitor
     pin_memory_fn
 
 

torchdata/datapipes/utils/__init__.py

@@ -10,4 +10,12 @@
 from torchdata.datapipes.utils.janitor import janitor
 from torchdata.datapipes.utils.pin_memory import pin_memory_fn
 
-__all__ = ["StreamWrapper", "janitor", "pin_memory_fn", "to_graph"]
+__all__ = [
+    "StreamWrapper",
+    "janitor",
+    "pin_memory_fn",
+    "to_graph",
+]
+
+# Please keep this list sorted
+assert __all__ == sorted(__all__)