Fix stubgen to add class docstrings and re-enable after version strin…

metaflow/cmd/develop/stub_generator.py

@@ -115,6 +115,11 @@ def __init__(self, output_dir: str, include_generated_for: bool = True):
         :type members_from_other_modules: List[str]
         """
 
+        # Let metaflow know we are in stubgen mode. This is sometimes useful to skip
+        # some processing like loading libraries, etc. It is used in Metaflow extensions
+        # so do not remove even if you do not see a use for it directly in the code.
+        os.environ["METAFLOW_STUBGEN"] = "1"
+
         self._write_generated_for = include_generated_for
         self._pending_modules = ["metaflow"]  # type: List[str]
         self._pending_modules.extend(get_aliased_modules())
@@ -398,6 +403,18 @@ def _generate_class_stub(self, name: str, clazz: type) -> str:
         name_with_module = self._get_element_name_with_module(clazz.__class__)
         buff.write("metaclass=" + name_with_module + "):\n")
 
+        # Add class docstring
+        if clazz.__doc__:
+            buff.write('%s"""\n' % TAB)
+            my_doc = cast(str, deindent_docstring(clazz.__doc__))
+            init_blank = True
+            for line in my_doc.split("\n"):
+                if init_blank and len(line.strip()) == 0:
+                    continue
+                init_blank = False
+                buff.write("%s%s\n" % (TAB, line.rstrip()))
+            buff.write('%s"""\n' % TAB)
+
         # For NamedTuple, we have __annotations__ but no __init__. In that case,
         # we are going to "create" a __init__ function with the annotations
         # to show what the class takes.

metaflow/cmd/develop/stubs.py

@@ -170,7 +170,7 @@ def install(ctx: Any, force: bool):
                     "Metaflow stubs are already installed and valid -- use --force to reinstall"
                 )
             return
-    mf_version, _ = get_mf_version()
+    mf_version, _ = get_mf_version(True)
     with tempfile.TemporaryDirectory() as tmp_dir:
         with open(os.path.join(tmp_dir, "setup.py"), "w") as f:
             f.write(
@@ -261,10 +261,10 @@ def split_version(vers: str) -> Tuple[str, Optional[str]]:
     return vers_split[0], vers_split[1]
 
 
-def get_mf_version() -> Tuple[str, Optional[str]]:
+def get_mf_version(public: bool = False) -> Tuple[str, Optional[str]]:
     from metaflow.metaflow_version import get_version
 
-    return split_version(get_version())
+    return split_version(get_version(public))
 
 
 def get_stubs_version(stubs_root_path: Optional[str]) -> Tuple[str, Optional[str]]:

metaflow/metaflow_version.py

@@ -195,11 +195,14 @@ def get_version(public=False):
         if ext_version is None:
             ext_version = getattr(extension_module, "__version__", "<unk>")
         # Update the package information about reported version for the extension
-        update_package_info(
-            package_name=pkg_name,
-            extension_name=ext_name,
-            package_version=ext_version,
-        )
+        # (only for the full info which is called at least once -- if we update more
+        # it will error out since we can only update_package_info once)
+        if not public:
+            update_package_info(
+                package_name=pkg_name,
+                extension_name=ext_name,
+                package_version=ext_version,
+            )
         ext_versions.append("%s(%s)" % (ext_name, ext_version))
 
     # We now have all the information about extensions so we can form the final string