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__ importannotations

importos.path
importre
importsubprocess
importsys
importtempfile

REVEAL_TYPE_START="reveal_type("
REVEAL_TYPE_END=")"


defupdate_line(line: str, s: str, pos: int) ->str:
returnline[:pos] +s+line[pos:]


defrun_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
 )
assertisinstance(
proc.stdout, bytes
 ) # Guaranteed to be true because we called run with universal_newlines=False
returnproc.stdout.decode(encoding="utf-8")


defget_revealed_type(line: str, relevant_file: str, relevant_line: int) ->str|None:
m=re.match(r'(.+?):(\d+): note: Revealed type is "(.*)"$', line)
ifmandint(m.group(2)) ==relevant_lineandos.path.samefile(relevant_file, m.group(1)):
returnm.group(3)
else:
returnNone


defprocess_output(output: str, filename: str, start_line: int) ->tuple[str|None, bool]:
error_found=False
forlineinoutput.splitlines():
t=get_revealed_type(line, filename, start_line)
ift:
returnt, error_found
elif"error:"inline:
error_found=True
returnNone, True# finding no reveal_type is an error


defmain() ->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)
withopen(filename) asf:
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)
withtempfile.NamedTemporaryFile(mode="w", prefix="mypy") astmp_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)
ifrevealed_type:
print(revealed_type)
iferror:
print(output)
exit(int(error))


if__name__=="__main__":
main()