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

…g change (#2051)

Two changes:
  - add docstrings to the class object
  - fix issue when generating stubs when Metaflow is not a released version

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