Remove layers from deep imports where possible

denest_imports.py

"""
Search for imports that are not at the top-level and replace them with
top-level imports.

from x.y import Z -> from x import Z

NOTE: doesn't handle comments or `import x as y`
"""

import argparse
import importlib
from collections import defaultdict
from functools import lru_cache
from typing import Tuple

import redbaron


def update_all_imports(contents: redbaron.RedBaron):
    for imp in contents.find_all("fromimport"):
        sources = [source.value for source in imp.value]
        final_mapping = defaultdict(set)
        for target in imp.targets:
            if target.target:
                # currently doesn't handle imports like `x as y`
                continue

            final_sources = promote_import(tuple(sources), target.value)
            final_mapping[tuple(final_sources)].add(target.value)

        if final_mapping != {tuple(sources): {t.value for t in imp.targets}}:
            for new_sources, new_targets in final_mapping.items():
                imp.insert_after(
                    "from {} import {}".format(
                        ".".join(new_sources),
                        ", ".join(new_targets),
                    )
                )
                imp.parent.remove(imp)


@lru_cache()
def promote_import(sources: Tuple, target: str):
    working_sources = list(sources)

    while working_sources:
        try:
            module = importlib.import_module(".".join(working_sources[:-1]))
            getattr(module, target)
        except (ValueError, ImportError, AttributeError):
            return working_sources
        else:
            working_sources.pop()

    return working_sources


def update_file(filename, *, dry_run):
    with open(filename, "r") as fp:
        baron = redbaron.RedBaron(fp.read())

    update_all_imports(baron)

    if dry_run:
        print(baron.dumps())
    else:
        with open(filename, "w") as fp:
            fp.write(baron.dumps())


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("files", nargs="+")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    for filename in args.files:
        update_file(filename, dry_run=args.dry_run)