Merge pull request #40 from pvandyken/feat/typed_lazy_load

Add type-checker compatible lazy import mode

README.rst

@@ -201,6 +201,14 @@ Although if you are willing to depend on the
 package and the ``--lazy_loader`` option (new as of 1.0.0), then this
 boilerplate is no longer needed.
 
+By default, lazy imports are not compatibly with statically typed projects (e.g
+using mypy or pyright), however, if the
+`lazy_loader <https://pypi.org/project/lazy_loader/>`_
+package is used, the ``--lazy_loader_typed`` option can be specified to generate
+``__init.pyi__`` files in addition to lazily evaulated ``__init.py__`` files.
+These interface files are understood by static type checkers and allow the
+combination of lazy loading with static type checking.
+
 
 Command Line Usage
 ------------------
@@ -267,6 +275,8 @@ Running ``mkint --help`` displays:
       --relative            Use relative . imports instead of <modname>
       --lazy                Use lazy imports with more boilerplate but no dependencies (Python >= 3.7 only!)
       --lazy_loader         Use lazy imports with less boilerplate but requires the lazy_loader module (Python >= 3.7 only!)
+      --lazy_loader_typed   Use lazy imports with the lazy_loader module, additionally generating 
+                            ``__init__.pyi`` files for static typing (e.g. with mypy or pyright) (Python >= 3.7 only!)
       --black               Use black formatting
       --lazy_boilerplate LAZY_BOILERPLATE
                             Code that defines a custom lazy_import callable

mkinit/__main__.py

@@ -111,6 +111,16 @@ def main():
         help="Use lazy imports with less boilerplate but requires the lazy_loader module (Python >= 3.7 only!)",
     )
 
+    lazy_group.add_argument(
+        "--lazy_loader_typed",
+        action="store_true",
+        default=False,
+        help=(
+            "Use lazy imports with the lazy_loader module, additionally generating "
+            "``__init__.pyi`` files for static typing (e.g. with mypy or pyright) (Python >= 3.7 only!)"
+        ),
+    )
+
     parser.add_argument(
         "--black",
         action="store_true",
@@ -147,14 +157,16 @@ def main():
     parser.add_argument("--version", action="store_true", help="print version and exit")
 
     import os
-    if os.environ.get('MKINIT_ARGPARSE_LOOSE', ''):
+
+    if os.environ.get("MKINIT_ARGPARSE_LOOSE", ""):
         args, unknown = parser.parse_known_args()
     else:
         args = parser.parse_args()
     ns = args.__dict__.copy()
 
     if ns["version"]:
         import mkinit
+
         print(mkinit.__version__)
         return
 
@@ -166,17 +178,30 @@ def main():
     verbose = ns["verbose"]
     dry = ns["dry"]
 
-    if ns['lazy_boilerplate'] and ns['lazy_loader']:
-        raise ValueError('--lazy_boilerplate cannot be specified with --lazy_loader. Use --lazy instead.')
+    if ns["lazy_boilerplate"] and (ns["lazy_loader"] or ns["lazy_loader_typed"]):
+        raise ValueError(
+            "--lazy_boilerplate cannot be specified with --lazy_loader or --lazy_loader_typed. Use --lazy instead."
+        )
+
+    if not ns["with_all"] and ns["lazy_loader_typed"]:
+        raise ValueError("--noall cannot be combined with --lazy_loader_typed")
+
+    if ns["lazy_loader_typed"] and not ns["relative"]:
+        print(
+            "WARNING: specifying --lazy-loader-typed implicitly enables --relative, as "
+            "`lazy-loader` stub support requires relative imports. (Explicitly specify "
+            "--relative to remove this warning.)"
+        )
 
     # Formatting options
     options = {
         "with_attrs": ns["with_attrs"],
         "with_mods": ns["with_mods"],
         "with_all": ns["with_all"],
-        "relative": ns["relative"],
+        "relative": ns["relative"] or ns["lazy_loader_typed"],
         "lazy_import": ns["lazy"],
-        "lazy_loader": ns["lazy_loader"],
+        "lazy_loader": ns["lazy_loader"] or ns["lazy_loader_typed"],
+        "lazy_loader_typed": ns["lazy_loader_typed"],
         "lazy_boilerplate": ns["lazy_boilerplate"],
         "use_black": ns["black"],
     }
@@ -202,8 +227,12 @@ def main():
 
     # print('ns = {!r}'.format(ns))
     static_mkinit.autogen_init(
-        modname_or_path, respect_all=respect_all, options=options, dry=dry,
-        diff=diff, recursive=ns['recursive'],
+        modname_or_path,
+        respect_all=respect_all,
+        options=options,
+        dry=dry,
+        diff=diff,
+        recursive=ns["recursive"],
     )
 
 

mkinit/formatting.py

@@ -35,6 +35,7 @@ def _ensure_options(given_options=None):
         "relative": False,
         "lazy_import": False,
         "lazy_loader": False,
+        "lazy_loader_typed": False,
         "lazy_boilerplate": None,
         "use_black": False,
     }
@@ -46,7 +47,7 @@ def _ensure_options(given_options=None):
     return options
 
 
-def _insert_autogen_text(modpath, initstr):
+def _insert_autogen_text(modpath, initstr, interface=False):
     """
     Creates new text for `__init__.py` containing the autogenerated code.
 
@@ -56,7 +57,7 @@ def _insert_autogen_text(modpath, initstr):
     """
 
     # Get path to init file so we can overwrite it
-    init_fpath = join(modpath, "__init__.py")
+    init_fpath = join(modpath, "__init__.pyi" if interface else "__init__.py")
     logger.debug("inserting initstr into: {!r}".format(init_fpath))
 
     if exists(init_fpath):
@@ -70,12 +71,11 @@ def _insert_autogen_text(modpath, initstr):
 
     QUICKFIX_REMOVE_LEADING_NEWLINES = 1
     if QUICKFIX_REMOVE_LEADING_NEWLINES:
-        initstr_ = initstr_.lstrip('\n')
+        initstr_ = initstr_.lstrip("\n")
 
     new_lines = lines[:startline] + [initstr_] + lines[endline:]
 
     new_text = "".join(new_lines).rstrip() + "\n"
-    print(new_text)
     return init_fpath, new_text
 
 
@@ -404,6 +404,10 @@ def append_part(new_part):
             """
         ).rstrip("\n")
         template = textwrap.dedent(
+            """
+            __getattr__, __dir__, __all__ = lazy_loader.attach_stub(__name__, __file__)
+            """
+            if options["lazy_loader_typed"] else
             """
             __getattr__, __dir__, __all__ = lazy_loader.attach(
                 __name__,

mkinit/static_mkinit.py

@@ -6,7 +6,7 @@
 from mkinit.util import util_import
 from mkinit.util.util_diff import difftext
 from mkinit.top_level_ast import TopLevelVisitor
-from mkinit.formatting import _initstr, _insert_autogen_text
+from mkinit.formatting import _initstr, _insert_autogen_text, _ensure_options
 from os.path import abspath
 from os.path import exists
 from os.path import join
@@ -83,6 +83,7 @@ def autogen_init(
     logger.info(
         "Autogenerating __init__ for modpath_or_name={}".format(modpath_or_name)
     )
+    options = _ensure_options(options)
     modpath = _rectify_to_modpath(modpath_or_name)
 
     if recursive:
@@ -106,33 +107,46 @@ def autogen_init(
                 diff=diff,
                 recursive=False,
             )
-
-    else:
-        initstr = static_init(
-            modpath, submodules=submodules, respect_all=respect_all, options=options
+        return
+
+    if options["lazy_loader_typed"] and options["lazy_loader"]:
+        autogen_init(
+            modpath,
+            submodules=None,
+            respect_all=respect_all,
+            options={**options, "lazy_loader": False},
+            dry=dry,
+            diff=diff,
+            recursive=False,
         )
-        init_fpath, new_text = _insert_autogen_text(modpath, initstr)
-        if dry:
-            logger.info("(DRY) would write updated file: %r" % init_fpath)
-            if diff:
-                # Display difference
-                try:
-                    with open(init_fpath, "r") as file:
-                        old_text = file.read()
-                except Exception:
-                    old_text = ""
-                display_text = difftext(
-                    old_text, new_text, colored=True, context_lines=3
-                )
-                print(display_text)
-            else:
-                print(new_text)
-            return init_fpath, new_text
+
+    initstr = static_init(
+        modpath, submodules=submodules, respect_all=respect_all, options=options
+    )
+    init_fpath, new_text = _insert_autogen_text(
+        modpath,
+        initstr,
+        interface=options["lazy_loader_typed"] and not options["lazy_loader"],
+    )
+    if dry:
+        logger.info("(DRY) would write updated file: %r" % init_fpath)
+        if diff:
+            # Display difference
+            try:
+                with open(init_fpath, "r") as file:
+                    old_text = file.read()
+            except Exception:
+                old_text = ""
+            display_text = difftext(old_text, new_text, colored=True, context_lines=3)
+            print(display_text)
         else:
-            logger.info("writing updated file: %r" % init_fpath)
-            # print(new_text)
-            with open(init_fpath, "w") as file_:
-                file_.write(new_text)
+            print(new_text)
+        return init_fpath, new_text
+    else:
+        logger.info("writing updated file: %r" % init_fpath)
+        # print(new_text)
+        with open(init_fpath, "w") as file_:
+            file_.write(new_text)
 
 
 def _rectify_to_modpath(modpath_or_name):

requirements/tests.txt

@@ -21,6 +21,8 @@ pytest-cov>=2.9.0           ; python_version < '3.6.0'  and python_version >= '3
 pytest-cov>=2.8.1           ; python_version < '3.5.0'  and python_version >= '3.4.0'   # Python 3.4
 pytest-cov>=2.8.1           ; python_version < '2.8.0'  and python_version >= '2.7.0'   # Python 2.7
 
+lazy-loader>=0.3            ;                               python_version >= '3.7.0'
+
 # xdev availpkg xdoctest
 # xdev availpkg coverage
 coverage>=6.1.1     ;                            python_version >= '3.10'    # Python 3.10+

tests/test_with_dummy.py

@@ -6,6 +6,7 @@
 import os
 from os.path import join
 from os.path import dirname
+import pytest
 try:
     from packaging.version import parse as LooseVersion
 except ImportError:
@@ -342,7 +343,12 @@ def test_dynamic_init():
         assert want in text, "missing {}".format(want)
 
 
-def test_lazy_import():
+@pytest.mark.parametrize(["option", "typed"], [
+    ("lazy_import", False),
+    ("lazy_loader", False),
+    ("lazy_loader", True),
+])
+def test_lazy_import(option, typed):
     """
     python ~/code/mkinit/tests/test_with_dummy.py test_lazy_import
     """
@@ -356,7 +362,12 @@ def test_lazy_import():
     paths = make_dummy_package(cache_dpath)
 
     pkg_path = paths["root"]
-    mkinit.autogen_init(pkg_path, options={"lazy_import": 1}, dry=False, recursive=True)
+    mkinit.autogen_init(
+        pkg_path,
+        options={option: 1, "lazy_loader_typed": typed, "relative": typed},
+        dry=False,
+        recursive=True
+    )
 
     if LooseVersion("{}.{}".format(*sys.version_info[0:2])) < LooseVersion("3.7"):
         pytest.skip()
@@ -375,11 +386,15 @@ def test_lazy_import():
         mkinit_dummy_module.a_very_nested_function()
 
 
-def test_recursive_lazy_autogen():
+@pytest.mark.parametrize(["option", "typed"], [
+    ("lazy_import", False),
+    ("lazy_loader", False),
+    ("lazy_loader", True),
+])
+def test_recursive_lazy_autogen(option, typed):
     """
     xdoctest ~/code/mkinit/tests/test_with_dummy.py test_recursive_lazy_autogen
     """
-    import pytest
     import mkinit
     import os
 
@@ -392,7 +407,12 @@ def test_recursive_lazy_autogen():
     )
     pkg_path = paths["root"]
 
-    mkinit.autogen_init(pkg_path, options={"lazy_import": 1}, dry=False, recursive=True)
+    mkinit.autogen_init(
+        pkg_path,
+        options={option: 1, "lazy_loader_typed": typed, "relative": typed},
+        dry=False,
+        recursive=True
+    )
 
     if LooseVersion("{}.{}".format(*sys.version_info[0:2])) < LooseVersion("3.7"):
         pytest.skip()
@@ -425,6 +445,38 @@ def test_recursive_lazy_autogen():
         )
         mkinit_rec_lazy_autogen.a_very_nested_function()
 
+def test_typed_pyi_file():
+    """
+    xdoctest ~/code/mkinit/tests/test_with_dummy.py test_recursive_lazy_autogen
+    """
+    import mkinit
+    import os
+
+    if sys.version_info[0:2] < (3, 7):
+        pytest.skip('Only 3.7+ has lazy imports')
+
+    cache_dpath = ub.Path.appdir("mkinit/tests").ensuredir()
+    paths = make_dummy_package(cache_dpath)
+    pkg_path = paths["root"]
+
+    mkinit.autogen_init(
+        pkg_path,
+        options={"lazy_loader": 1, "lazy_loader_typed": True, "relative": True},
+        dry=False,
+        recursive=True
+    )
+    if LooseVersion("{}.{}".format(*sys.version_info[0:2])) < LooseVersion("3.7"):
+        pytest.skip()
+
+    dpath = dirname(paths["root"])
+    with ub.util_import.PythonPathContext(os.fspath(dpath)):
+        import mkinit_dummy_module
+        text = ub.Path(*mkinit_dummy_module.__path__, "__init__.pyi").read_text()
+        assert "from . import avery" in text
+        assert "from . import subdir1" in text
+        assert "from .submod1 import (attr1, attr2" in text
+        assert "__all__ = ['a_very_nested_function'" in text
+
 
 def test_recursive_eager_autogen():
     """