diff --git a/tools/dev/generate_pyi_files.py b/tools/dev/generate_pyi_files.py deleted file mode 100644 index c6ae36e7dd4..00000000000 --- a/tools/dev/generate_pyi_files.py +++ /dev/null @@ -1,81 +0,0 @@ -# Authors: The MNE-Python contributors. -# License: BSD-3-Clause -# Copyright the MNE-Python contributors. - -import os -import sys -from importlib import import_module -from pathlib import Path - -import ast_comments as ast -import black - -import mne - - -class RewriteAssign(ast.NodeTransformer): - """NodeTransformer to replace lazy attach with attach_stub.""" - - def visit_Assign(self, node): - """Replace lazy attach assignment with stub assignment.""" - if not hasattr(node.targets[0], "dims"): - return node - - ids = [name.id for name in node.targets[0].dims] - if ids == ["__getattr__", "__dir__", "__all__"]: - return ast.parse( - "__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)\n" - ) - return node - - -pyi_mode = black.Mode(is_pyi=True) -root = Path(mne.__file__).parent -inits = root.rglob("__init__.py") - -for init in inits: - # skip init files that don't lazy load (e.g., tests) - code = init.read_text("utf-8") - if "import lazy_loader as lazy" not in code: - continue - # get the AST - tree = ast.parse(code) - nodes = [node for node in tree.body if isinstance(node, ast.Assign)] - assert len(nodes) == 1 - node = nodes[0] - keywords = node.value.keywords - # get submodules - import_lines = list() - assert keywords[0].arg == "submodules" - # for submod in keywords[0].value.elts: - # import_lines.append(f"import {submod.value}") - submods = [submod.value for submod in keywords[0].value.elts] - if len(submods): - import_lines.append(f"from . import {', '.join(submods)}") - # get attrs - assert keywords[1].arg == "submod_attrs" - _dict = keywords[1].value - for key, vals in zip(_dict.keys, _dict.values): - attrs = [attr.value for attr in vals.elts] - import_lines.append(f"from .{key.value} import {', '.join(attrs)}") - # format - import_lines = black.format_str("\n".join(import_lines), mode=pyi_mode) - # get __all__ - import_path = str(init.parent.relative_to(root.parent)).replace(os.sep, ".") - import_module(import_path) - _all = black.format_str( - f"__all__ = {repr(sys.modules[import_path].__all__)}\n", - mode=pyi_mode, - ) - # write __init__.pyi - outfile = init.with_suffix(".pyi") - with open(outfile, "w") as fid: - fid.write(_all) - fid.write(import_lines) - # rewrite __init__.py - new_tree = RewriteAssign().visit(tree) - new_tree = ast.fix_missing_locations(new_tree) - new_code = ast.unparse(new_tree) - formatted_code = black.format_str(new_code, mode=black.Mode()) - with open(init, "w") as fid: - fid.write(formatted_code)