autoimportext.py

#!/usr/bin/env python3

import ast
import os
import sys
import contextlib
import tempfile
from pathlib import Path
from typing import Dict, Sequence

import toml
from autoimport.entrypoints import cli
from maison.config import ProjectConfig


def gather_import_statements() -> str:
    raise Exception("TODO")


def _parse_import_config(config_code: str) -> Dict[str, str]:
    ns: Dict[str, str] = {}
    # An alternative approach would be to `exec(config_code, ns)`,
    # which would be able to support `if`s, wildcard imports, and other dynamicity,
    # but would not be able to support anything but classes/functions.
    statements = ast.parse(config_code).body
    for stmt in statements:
        if isinstance(stmt, ast.Import):
            for imp in stmt.names:
                if imp.asname:
                    ns[imp.asname] = f"import {imp.name} as {imp.asname}"
                else:
                    ns[imp.name] = f"import {imp.name}"
        elif isinstance(stmt, ast.ImportFrom):
            for imp in stmt.names:
                if imp.asname:
                    ns[
                        imp.asname
                    ] = f"from {stmt.module} import {imp.name} as {imp.asname}"
                else:
                    ns[imp.name] = f"from {stmt.module} import {imp.name}"

    return ns


def _parse_all_import_configs(
    config_paths: Sequence[str] = ("~/.autoimport_ns.py", "./.autoimport_ns.py"),
) -> Dict[str, str]:
    ns: Dict[str, str] = {}
    for config_path in config_paths:
        config_path = os.path.expanduser(config_path)
        try:
            with open(config_path) as fobj:
                config_code = fobj.read()
        except FileNotFoundError:
            continue
        ns.update(_parse_import_config(config_code))
    return ns

def _parse_project_config() -> Dict[str, str]:
    return ProjectConfig(project_name="autoimport").to_dict()


def _make_full_config() -> str:
    """
    WARNING: changes the current directory.
    """
    project_config = _parse_project_config()
    imports_config = _parse_all_import_configs()
    full_config = {
        "tool": {
            "autoimport": {
                **project_config,
                "common_statements": {
                    **(project_config.get("common_statements") or {}),
                    **imports_config,
                },
            },
        },
    }
    return toml.dumps(full_config)


def _normalize_arg(maybe_path: str) -> str:
    if maybe_path.startswith("-"):
        return maybe_path
    return str(Path(maybe_path).resolve())


def _main():
    full_config = _make_full_config()
    norm_paths = [_normalize_arg(arg) for arg in sys.argv[1:]]
    if not norm_paths:
        norm_paths = [str(path.resolve()) for path in Path.cwd().rglob('**/*.py')]

    sys.argv[1:] = norm_paths
    with tempfile.TemporaryDirectory() as tempdir:
        os.chdir(tempdir)
        with open("pyproject.toml", "w") as fobj:
            fobj.write(full_config)
        print(full_config[:1000])
        return cli.cli()


if __name__ == "__main__":
    _main()