add_future_annotations_import.py

import re
import subprocess
import sys
from pathlib import Path

FUTURE_ANNOTATIONS_IMPORT_REGEX = (
    r"^from __future__ import (?:.*, )?annotations(?:, .*)?"
)


def add_future_annotations_import(file: Path, reformat: bool = True) -> None:
    file_str = str(file)
    source = file.read_text()
    if "/tests/" in file_str or "/migrations/" in file_str:
        print(f"🤷 {file}")
        return
    if re.findall(FUTURE_ANNOTATIONS_IMPORT_REGEX, source, re.MULTILINE):
        print(f"✅ {file}")
        return
    print(f"❌ {file}")
    lines = source.splitlines()
    in_docstring = False
    for line_index, line in enumerate(lines):
        if line.startswith("#"):
            continue
        if line == '"""':
            in_docstring = not in_docstring
            continue
        else:
            if line.startswith('"""'):
                in_docstring = True
            if line.endswith('"""'):
                in_docstring = False
                continue
        if in_docstring:
            continue
        lines.insert(
            line_index, "from __future__ import annotations",
        )
        break
    source = "\n".join(lines)
    source = f"{source}\n"
    file.write_text(source)

    if reformat:
        subprocess.run(
            ["isort", file_str], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        subprocess.run(
            ["black", file_str], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )


def main() -> None:
    if len(sys.argv) != 2:
        print(f"Usage: python3 {sys.argv[0]} PROJECT_DIR")
        sys.exit(1)

    for file in Path(sys.argv[1]).glob("yunojuno/**/*.py"):
        add_future_annotations_import(file)


if __name__ == "__main__":
    main()