use the top-level SDK path (e.g. airflow.sdk.DAG)

Author: sunank200

task-sdk/docs/api.rst

@@ -69,7 +69,7 @@ Assets
 Everything else
 ---------------
 
-.. autoapimodule:: airflow.sdk
+.. autoapi_module:: airflow.sdk
   :members:
   :exclude-members: BaseOperator, DAG, dag, asset, Asset, AssetAlias, AssetAll, AssetAny, AssetWatcher, TaskGroup, XComArg, get_current_context, get_parsing_context
   :undoc-members:

task-sdk/docs/conf.py

@@ -43,6 +43,7 @@
 autoapi_options = [
     "undoc-members",
     "members",
+    "imported-members",
 ]
 autoapi_add_toctree_entry = False
 autoapi_generate_api_docs = False
@@ -78,6 +79,6 @@ def skip_util_classes(app, objtype, name, obj, skip, options):
     return skip
 
 
-def setup(sphinx):
-    # sphinx.connect("autoapi-skip-member", skip_util_classes)
-    ...
+def setup(app):
+    # Skip utility classes from definitions modules, but expose DAG under airflow.sdk.DAG
+    app.connect("autoapi-skip-member", skip_util_classes)

task-sdk/tests/test_docs_inventory.py

@@ -16,6 +16,11 @@
 # under the License.
 from __future__ import annotations
 
+import sys
+from pathlib import Path
+
+# Add the SDK src directory to sys.path so that importlib loads our airflow.sdk module
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 import importlib
 import shutil
 import subprocess
@@ -62,7 +67,11 @@ def test_docs_inventory_matches_public_api(tmp_path):
     assert inv_path.exists(), "objects.inv not found after docs build"
 
     inv = read_inventory(inv_path)
-    documented = {name.rsplit(".", 1)[-1] for name in inv.keys()}
+    documented = {
+        name.rsplit(".", 1)[-1]
+        for name in inv.keys()
+        if name.startswith("airflow.sdk.") and name.count(".") == 2
+    }
     sdk = importlib.import_module("airflow.sdk")
     public = set(getattr(sdk, "__all__", []))