From 0d6bbc02ed6a50000d3779385d121a9d48d36702 Mon Sep 17 00:00:00 2001 From: Romain Cledat Date: Sat, 21 Sep 2024 14:03:28 -0700 Subject: [PATCH] Fix stubgen to add class docstrings and re-enable after version string change 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 | 17 +++++++++++++++++ metaflow/cmd/develop/stubs.py | 6 +++--- metaflow/metaflow_version.py | 13 ++++++++----- 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/metaflow/cmd/develop/stub_generator.py b/metaflow/cmd/develop/stub_generator.py index 34491fedb17..0e6e9f86ae6 100644 --- a/metaflow/cmd/develop/stub_generator.py +++ b/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. diff --git a/metaflow/cmd/develop/stubs.py b/metaflow/cmd/develop/stubs.py index 49c00486c3f..64bdc7b5a16 100644 --- a/metaflow/cmd/develop/stubs.py +++ b/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]]: diff --git a/metaflow/metaflow_version.py b/metaflow/metaflow_version.py index 139948e3b80..e3be8ed7956 100644 --- a/metaflow/metaflow_version.py +++ b/metaflow/metaflow_version.py @@ -195,11 +195,14 @@ def get_version(public=False): if ext_version is None: ext_version = getattr(extension_module, "__version__", "") # 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