infra: create api rst for specific pkg

docs/api_reference/create_api_rst.py

@@ -344,28 +344,26 @@ def _doc_first_line(package_name: str) -> str:
     return f".. {package_name.replace('-', '_')}_api_reference:\n\n"
 
 
-def main() -> None:
+def main(dirs: Optional[list] = None) -> None:
     """Generate the api_reference.rst file for each package."""
     print("Starting to build API reference files.")
-    for dir in os.listdir(ROOT_DIR / "libs"):
+    if not dirs:
+        dirs = [dir_ for dir_ in os.listdir(ROOT_DIR / "libs") if dir_ not in ("cli", "partners")]
+        dirs += os.listdir(ROOT_DIR / "libs" / "partners")
+    for dir_ in dirs:
         # Skip any hidden directories
         # Some of these could be present by mistake in the code base
         # e.g., .pytest_cache from running tests from the wrong location.
-        if dir.startswith("."):
-            print("Skipping dir:", dir)
-            continue
-
-        if dir in ("cli", "partners"):
+        if dir_.startswith("."):
+            print("Skipping dir:", dir_)
             continue
         else:
-            print("Building package:", dir)
-            _build_rst_file(package_name=dir)
-    partner_packages = os.listdir(ROOT_DIR / "libs" / "partners")
-    print("Building partner packages:", partner_packages)
-    for dir in partner_packages:
-        _build_rst_file(package_name=dir)
+            print("Building package:", dir_)
+            _build_rst_file(package_name=dir_)
     print("API reference files built.")
 
 
 if __name__ == "__main__":
-    main()
+    import sys
+    dirs = sys.argv[1:] or None
+    main(dirs=dirs)