Generate stubs properly for extensions with lazy loaded modules (#1765)

Author: romain-intel

metaflow/cmd/develop/stub_generator.py

@@ -30,6 +30,7 @@
 from metaflow import FlowSpec, step
 from metaflow.debug import debug
 from metaflow.decorators import Decorator, FlowDecorator
+from metaflow.extension_support import get_aliased_modules
 from metaflow.graph import deindent_docstring
 from metaflow.metaflow_version import get_version
 
@@ -116,6 +117,7 @@ def __init__(self, output_dir: str, include_generated_for: bool = True):
 
         self._write_generated_for = include_generated_for
         self._pending_modules = ["metaflow"]  # type: List[str]
+        self._pending_modules.extend(get_aliased_modules())
         self._root_module = "metaflow."
         self._safe_modules = ["metaflow.", "metaflow_extensions."]
 

metaflow/cmd/develop/stubs.py

@@ -323,17 +323,26 @@ def get_packages_for_stubs() -> Tuple[List[Tuple[str, str]], List[str]]:
         return [], []
 
     dist_list = []
-    for dist in _metadata_package.distributions():
+
+    # We check the type because if the user has multiple importlib metadata, for
+    # some reason it shows up multiple times.
+    interesting_dists = [
+        d
+        for d in _metadata_package.distributions()
         if any(
             [
-                pkg == "metaflow-stubs"
-                for pkg in (dist.read_text("top_level.txt") or "").split()
+                p == "metaflow-stubs"
+                for p in (d.read_text("top_level.txt") or "").split()
             ]
-        ):
-            # This is a package we care about
-            root_path = dist.locate_file("metaflow-stubs").as_posix()
-            dist_list.append((dist.metadata["Name"], root_path))
-            all_paths.discard(root_path)
+        )
+        and isinstance(d, _metadata_package.PathDistribution)
+    ]
+
+    for dist in interesting_dists:
+        # This is a package we care about
+        root_path = dist.locate_file("metaflow-stubs").as_posix()
+        dist_list.append((dist.metadata["Name"], root_path))
+        all_paths.discard(root_path)
     return dist_list, list(all_paths)