xref: /llvm-project/clang-tools-extra/clang-include-fixer/tool/clang-include-fixer.py (revision dd3c26a045c081620375a878159f536758baba6e)
1# This file is a minimal clang-include-fixer vim-integration. To install:
2# - Change 'binary' if clang-include-fixer is not on the path (see below).
3# - Add to your .vimrc:
4#
5#   noremap <leader>cf :pyf path/to/llvm/source/tools/clang/tools/extra/clang-include-fixer/tool/clang-include-fixer.py<cr>
6#
7# This enables clang-include-fixer for NORMAL and VISUAL mode. Change
8# "<leader>cf" to another binding if you need clang-include-fixer on a
9# different key.
10#
11# To set up clang-include-fixer, see
12# http://clang.llvm.org/extra/clang-include-fixer.html
13#
14# With this integration you can press the bound key and clang-include-fixer will
15# be run on the current buffer.
16#
17# It operates on the current, potentially unsaved buffer and does not create
18# or save any files. To revert a fix, just undo.
19
20from __future__ import print_function
21import argparse
22import difflib
23import json
24import re
25import subprocess
26import vim
27
28# set g:clang_include_fixer_path to the path to clang-include-fixer if it is not
29# on the path.
30# Change this to the full path if clang-include-fixer is not on the path.
31binary = "clang-include-fixer"
32if vim.eval('exists("g:clang_include_fixer_path")') == "1":
33    binary = vim.eval("g:clang_include_fixer_path")
34
35maximum_suggested_headers = 3
36if vim.eval('exists("g:clang_include_fixer_maximum_suggested_headers")') == "1":
37    maximum_suggested_headers = max(
38        1, vim.eval("g:clang_include_fixer_maximum_suggested_headers")
39    )
40
41increment_num = 5
42if vim.eval('exists("g:clang_include_fixer_increment_num")') == "1":
43    increment_num = max(1, vim.eval("g:clang_include_fixer_increment_num"))
44
45jump_to_include = False
46if vim.eval('exists("g:clang_include_fixer_jump_to_include")') == "1":
47    jump_to_include = vim.eval("g:clang_include_fixer_jump_to_include") != "0"
48
49query_mode = False
50if vim.eval('exists("g:clang_include_fixer_query_mode")') == "1":
51    query_mode = vim.eval("g:clang_include_fixer_query_mode") != "0"
52
53
54def GetUserSelection(message, headers, maximum_suggested_headers):
55    eval_message = message + "\n"
56    for idx, header in enumerate(headers[0:maximum_suggested_headers]):
57        eval_message += "({0}). {1}\n".format(idx + 1, header)
58    eval_message += "Enter (q) to quit;"
59    if maximum_suggested_headers < len(headers):
60        eval_message += " (m) to show {0} more candidates.".format(
61            min(increment_num, len(headers) - maximum_suggested_headers)
62        )
63
64    eval_message += "\nSelect (default 1): "
65    res = vim.eval("input('{0}')".format(eval_message))
66    if res == "":
67        # choose the top ranked header by default
68        idx = 1
69    elif res == "q":
70        raise Exception("   Insertion cancelled...")
71    elif res == "m":
72        return GetUserSelection(
73            message, headers, maximum_suggested_headers + increment_num
74        )
75    else:
76        try:
77            idx = int(res)
78            if idx <= 0 or idx > len(headers):
79                raise Exception()
80        except Exception:
81            # Show a new prompt on invalid option instead of aborting so that users
82            # don't need to wait for another clang-include-fixer run.
83            print("Invalid option: {}".format(res), file=sys.stderr)
84            return GetUserSelection(message, headers, maximum_suggested_headers)
85    return headers[idx - 1]
86
87
88def execute(command, text):
89    # Avoid flashing a cmd prompt on Windows.
90    startupinfo = None
91    if sys.platform.startswith("win32"):
92        startupinfo = subprocess.STARTUPINFO()
93        startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
94        startupinfo.wShowWindow = subprocess.SW_HIDE
95
96    p = subprocess.Popen(
97        command,
98        stdout=subprocess.PIPE,
99        stderr=subprocess.PIPE,
100        stdin=subprocess.PIPE,
101        startupinfo=startupinfo,
102    )
103    return p.communicate(input=text.encode("utf-8"))
104
105
106def InsertHeaderToVimBuffer(header, text):
107    command = [
108        binary,
109        "-stdin",
110        "-insert-header=" + json.dumps(header),
111        vim.current.buffer.name,
112    ]
113    stdout, stderr = execute(command, text)
114    if stderr:
115        raise Exception(stderr)
116    if stdout:
117        lines = stdout.splitlines()
118        sequence = difflib.SequenceMatcher(None, vim.current.buffer, lines)
119        line_num = None
120        for op in reversed(sequence.get_opcodes()):
121            if op[0] != "equal":
122                vim.current.buffer[op[1] : op[2]] = lines[op[3] : op[4]]
123            if op[0] == "insert":
124                # line_num in vim is 1-based.
125                line_num = op[1] + 1
126
127        if jump_to_include and line_num:
128            vim.current.window.cursor = (line_num, 0)
129
130
131# The vim internal implementation (expand("cword"/"cWORD")) doesn't support
132# our use case very well, we re-implement our own one.
133def get_symbol_under_cursor():
134    line = vim.eval('line(".")')
135    # column number in vim is 1-based.
136    col = int(vim.eval('col(".")')) - 1
137    line_text = vim.eval("getline({0})".format(line))
138    if len(line_text) == 0:
139        return ""
140    symbol_pos_begin = col
141    p = re.compile("[a-zA-Z0-9:_]")
142    while symbol_pos_begin >= 0 and p.match(line_text[symbol_pos_begin]):
143        symbol_pos_begin -= 1
144
145    symbol_pos_end = col
146    while symbol_pos_end < len(line_text) and p.match(line_text[symbol_pos_end]):
147        symbol_pos_end += 1
148    return line_text[symbol_pos_begin + 1 : symbol_pos_end]
149
150
151def main():
152    parser = argparse.ArgumentParser(
153        description="Vim integration for clang-include-fixer"
154    )
155    parser.add_argument("-db", default="yaml", help="clang-include-fixer input format.")
156    parser.add_argument("-input", default="", help="String to initialize the database.")
157    # Don't throw exception when parsing unknown arguments to make the script
158    # work in neovim.
159    # Neovim (at least v0.2.1) somehow mangles the sys.argv in a weird way: it
160    # will pass additional arguments (e.g. "-c script_host.py") to sys.argv,
161    # which makes the script fail.
162    args, _ = parser.parse_known_args()
163
164    # Get the current text.
165    buf = vim.current.buffer
166    text = "\n".join(buf)
167
168    if query_mode:
169        symbol = get_symbol_under_cursor()
170        if len(symbol) == 0:
171            print("Skip querying empty symbol.")
172            return
173        command = [
174            binary,
175            "-stdin",
176            "-query-symbol=" + get_symbol_under_cursor(),
177            "-db=" + args.db,
178            "-input=" + args.input,
179            vim.current.buffer.name,
180        ]
181    else:
182        # Run command to get all headers.
183        command = [
184            binary,
185            "-stdin",
186            "-output-headers",
187            "-db=" + args.db,
188            "-input=" + args.input,
189            vim.current.buffer.name,
190        ]
191    stdout, stderr = execute(command, text)
192    if stderr:
193        print(
194            "Error while running clang-include-fixer: {}".format(stderr),
195            file=sys.stderr,
196        )
197        return
198
199    include_fixer_context = json.loads(stdout)
200    query_symbol_infos = include_fixer_context["QuerySymbolInfos"]
201    if not query_symbol_infos:
202        print("The file is fine, no need to add a header.")
203        return
204    symbol = query_symbol_infos[0]["RawIdentifier"]
205    # The header_infos is already sorted by clang-include-fixer.
206    header_infos = include_fixer_context["HeaderInfos"]
207    # Deduplicate headers while keeping the order, so that the same header would
208    # not be suggested twice.
209    unique_headers = []
210    seen = set()
211    for header_info in header_infos:
212        header = header_info["Header"]
213        if header not in seen:
214            seen.add(header)
215            unique_headers.append(header)
216
217    if not unique_headers:
218        print("Couldn't find a header for {0}.".format(symbol))
219        return
220
221    try:
222        selected = unique_headers[0]
223        inserted_header_infos = header_infos
224        if len(unique_headers) > 1:
225            selected = GetUserSelection(
226                "choose a header file for {0}.".format(symbol),
227                unique_headers,
228                maximum_suggested_headers,
229            )
230            inserted_header_infos = [
231                header for header in header_infos if header["Header"] == selected
232            ]
233        include_fixer_context["HeaderInfos"] = inserted_header_infos
234
235        InsertHeaderToVimBuffer(include_fixer_context, text)
236        print("Added #include {0} for {1}.".format(selected, symbol))
237    except Exception as error:
238        print(error, file=sys.stderr)
239    return
240
241
242if __name__ == "__main__":
243    main()
244