misc/find_type.py

#!/usr/bin/env python3
# Usage: find_type.py FILENAME START_LINE START_COL END_LINE END_COL MYPY_AND_ARGS
# Prints out the type of the expression in the given location if the mypy run
# succeeds cleanly.  Otherwise, prints out the errors encountered.
# Note: this only works on expressions, and not assignment targets.
# Note: MYPY_AND_ARGS is should be the remainder of argv, not a single
# spaces-included argument.
# NOTE: Line numbers are 1-based; column numbers are 0-based.
#
#
# Example vim usage:
# function RevealType()
#   " Set this to the command you use to run mypy on your project.  Include the mypy invocation.
#   let mypycmd = 'python3 -m mypy mypy --incremental'
#   let [startline, startcol] = getpos("'<")[1:2]
#   let [endline, endcol] = getpos("'>")[1:2]
#   " Convert to 0-based column offsets
#   let startcol = startcol - 1
#   " Change this line to point to the find_type.py script.
#   execute '!python3 /path/to/mypy/misc/find_type.py % ' . startline . ' ' . startcol . ' ' . endline . ' ' . endcol . ' ' . mypycmd
# endfunction
# vnoremap <Leader>t :call RevealType()<CR>
#
# For an Emacs example, see misc/macs.el.

from __future__ import annotations

import os.path
import re
import subprocess
import sys
import tempfile

REVEAL_TYPE_START = "reveal_type("
REVEAL_TYPE_END = ")"


def update_line(line: str, s: str, pos: int) -> str:
    return line[:pos] + s + line[pos:]


def run_mypy(mypy_and_args: list[str], filename: str, tmp_name: str) -> str:
    proc = subprocess.run(
        mypy_and_args + ["--shadow-file", filename, tmp_name], stdout=subprocess.PIPE
    )
    assert isinstance(
        proc.stdout, bytes
    )  # Guaranteed to be true because we called run with universal_newlines=False
    return proc.stdout.decode(encoding="utf-8")


def get_revealed_type(line: str, relevant_file: str, relevant_line: int) -> str | None:
    m = re.match(r'(.+?):(\d+): note: Revealed type is "(.*)"$', line)
    if m and int(m.group(2)) == relevant_line and os.path.samefile(relevant_file, m.group(1)):
        return m.group(3)
    else:
        return None


def process_output(output: str, filename: str, start_line: int) -> tuple[str | None, bool]:
    error_found = False
    for line in output.splitlines():
        t = get_revealed_type(line, filename, start_line)
        if t:
            return t, error_found
        elif "error:" in line:
            error_found = True
    return None, True  # finding no reveal_type is an error


def main() -> None:
    filename, start_line_str, start_col_str, end_line_str, end_col_str, *mypy_and_args = sys.argv[
        1:
    ]
    start_line = int(start_line_str)
    start_col = int(start_col_str)
    end_line = int(end_line_str)
    end_col = int(end_col_str)
    with open(filename) as f:
        lines = f.readlines()
        lines[end_line - 1] = update_line(
            lines[end_line - 1], REVEAL_TYPE_END, end_col
        )  # insert after end_col
        lines[start_line - 1] = update_line(lines[start_line - 1], REVEAL_TYPE_START, start_col)
        with tempfile.NamedTemporaryFile(mode="w", prefix="mypy") as tmp_f:
            tmp_f.writelines(lines)
            tmp_f.flush()

            output = run_mypy(mypy_and_args, filename, tmp_f.name)
            revealed_type, error = process_output(output, filename, start_line)
            if revealed_type:
                print(revealed_type)
            if error:
                print(output)
            exit(int(error))


if __name__ == "__main__":
    main()