1import logging
2import tkinter as tk
3
4from thonny import get_workbench, jedi_utils
5
6logger = logging.getLogger(__name__)
7
8tree = None
9
10
11class BaseNameHighlighter:
12    def __init__(self, text):
13        self.text = text
14        self._update_scheduled = False
15
16    def get_positions_for(self, source, line, column):
17        raise NotImplementedError()
18
19    def get_positions(self):
20        index = self.text.index("insert")
21
22        # ignore if cursor in open string
23        if self.text.tag_prevrange("open_string", index) or self.text.tag_prevrange(
24            "open_string3", index
25        ):
26
27            return set()
28
29        source = self.text.get("1.0", "end")
30        index_parts = index.split(".")
31        line, column = int(index_parts[0]), int(index_parts[1])
32
33        return self.get_positions_for(source, line, column)
34
35    def schedule_update(self):
36        def perform_update():
37            try:
38                self.update()
39            finally:
40                self._update_scheduled = False
41
42        if not self._update_scheduled:
43            self._update_scheduled = True
44            self.text.after_idle(perform_update)
45
46    def update(self):
47        self.text.tag_remove("matched_name", "1.0", "end")
48
49        if get_workbench().get_option("view.name_highlighting") and self.text.is_python_text():
50            try:
51                positions = self.get_positions()
52                if len(positions) > 1:
53                    for pos in positions:
54                        start_index, end_index = pos[0], pos[1]
55                        self.text.tag_add("matched_name", start_index, end_index)
56            except Exception as e:
57                logger.exception("Problem when updating name highlighting", exc_info=e)
58
59
60class VariablesHighlighter(BaseNameHighlighter):
61    """This is heavy, but more correct solution for variables, than Script.usages provides
62    (at least for Jedi 0.10)"""
63
64    def _is_name_function_call_name(self, name):
65        stmt = name.get_definition()
66        return stmt.type == "power" and stmt.children[0] == name
67
68    def _is_name_function_definition(self, name):
69        scope = name.get_definition()
70        return (
71            isinstance(scope, tree.Function)
72            and hasattr(scope.children[1], "value")
73            and scope.children[1].value == name.value
74        )
75
76    def _get_def_from_function_params(self, func_node, name):
77        params = func_node.get_params()
78        for param in params:
79            if param.children[0].value == name.value:
80                return param.children[0]
81        return None
82
83    # copied from jedi's tree.py with a few modifications
84    def _get_statement_for_position(self, node, pos):
85        for c in node.children:
86            # sorted here, because the end_pos property depends on the last child having the last position,
87            # there seems to be a problem with jedi, where the children of a node are not always in the right order
88            if isinstance(c, tree.Class):
89                c.children.sort(key=lambda x: x.end_pos)
90            if c.start_pos <= pos <= c.end_pos:
91                if c.type not in ("decorated", "simple_stmt", "suite") and not isinstance(
92                    c, (tree.Flow, tree.ClassOrFunc)
93                ):
94                    return c
95                else:
96                    try:
97                        return jedi_utils.get_statement_of_position(c, pos)
98                    except AttributeError as e:
99                        logger.exception("Could not get statement of position", exc_info=e)
100        return None
101
102    def _is_global_stmt_with_name(self, node, name_str):
103        return (
104            isinstance(node, tree.BaseNode)
105            and node.type == "simple_stmt"
106            and isinstance(node.children[0], tree.GlobalStmt)
107            and node.children[0].children[1].value == name_str
108        )
109
110    def _find_definition(self, scope, name):
111        from jedi import parser_utils
112
113        # if the name is the name of a function definition
114        if isinstance(scope, tree.Function):
115            if scope.children[1] == name:
116                return scope.children[1]  # 0th child is keyword "def", 1st is name
117            else:
118                definition = self._get_def_from_function_params(scope, name)
119                if definition:
120                    return definition
121
122        for c in scope.children:
123            if (
124                isinstance(c, tree.BaseNode)
125                and c.type == "simple_stmt"
126                and isinstance(c.children[0], tree.ImportName)
127            ):
128                for n in c.children[0].get_defined_names():
129                    if n.value == name.value:
130                        return n
131                # print(c.path_for_name(name.value))
132            if (
133                isinstance(c, tree.Function)
134                and c.children[1].value == name.value
135                and not isinstance(parser_utils.get_parent_scope(c), tree.Class)
136            ):
137                return c.children[1]
138            if isinstance(c, tree.BaseNode) and c.type == "suite":
139                for x in c.children:
140                    if self._is_global_stmt_with_name(x, name.value):
141                        return self._find_definition(parser_utils.get_parent_scope(scope), name)
142                    if isinstance(x, tree.Name) and x.is_definition() and x.value == name.value:
143                        return x
144                    def_candidate = self._find_def_in_simple_node(x, name)
145                    if def_candidate:
146                        return def_candidate
147
148        if not isinstance(scope, tree.Module):
149            return self._find_definition(parser_utils.get_parent_scope(scope), name)
150
151        # if name itself is the left side of an assignment statement, then the name is the definition
152        if name.is_definition():
153            return name
154
155        return None
156
157    def _find_def_in_simple_node(self, node, name):
158        if isinstance(node, tree.Name) and node.is_definition() and node.value == name.value:
159            return name
160        if not isinstance(node, tree.BaseNode):
161            return None
162        for c in node.children:
163            return self._find_def_in_simple_node(c, name)
164
165    def _get_dot_names(self, stmt):
166        try:
167            if (
168                hasattr(stmt, "children")
169                and len(stmt.children) >= 2
170                and hasattr(stmt.children[1], "children")
171                and len(stmt.children[1].children) >= 1
172                and hasattr(stmt.children[1].children[0], "value")
173                and stmt.children[1].children[0].value == "."
174            ):
175                return stmt.children[0], stmt.children[1].children[1]
176        except Exception as e:
177            logger.exception("_get_dot_names", exc_info=e)
178            return ()
179        return ()
180
181    def _find_usages(self, name, stmt):
182        from jedi import parser_utils
183
184        # check if stmt is dot qualified, disregard these
185        dot_names = self._get_dot_names(stmt)
186        if len(dot_names) > 1 and dot_names[1].value == name.value:
187            return set()
188
189        # search for definition
190        definition = self._find_definition(parser_utils.get_parent_scope(name), name)
191
192        searched_scopes = set()
193
194        is_function_definition = (
195            self._is_name_function_definition(definition) if definition else False
196        )
197
198        def find_usages_in_node(node, global_encountered=False):
199            names = []
200            if isinstance(node, tree.BaseNode):
201                if parser_utils.is_scope(node):
202                    global_encountered = False
203                    if node in searched_scopes:
204                        return names
205                    searched_scopes.add(node)
206                    if isinstance(node, tree.Function):
207                        d = self._get_def_from_function_params(node, name)
208                        if d and d != definition:
209                            return []
210
211                for c in node.children:
212                    dot_names = self._get_dot_names(c)
213                    if len(dot_names) > 1 and dot_names[1].value == name.value:
214                        continue
215                    sub_result = find_usages_in_node(c, global_encountered=global_encountered)
216
217                    if sub_result is None:
218                        if not parser_utils.is_scope(node):
219                            return (
220                                None
221                                if definition and node != parser_utils.get_parent_scope(definition)
222                                else [definition]
223                            )
224                        else:
225                            sub_result = []
226                    names.extend(sub_result)
227                    if self._is_global_stmt_with_name(c, name.value):
228                        global_encountered = True
229            elif isinstance(node, tree.Name) and node.value == name.value:
230                if definition and definition != node:
231                    if self._is_name_function_definition(node):
232                        if isinstance(
233                            parser_utils.get_parent_scope(parser_utils.get_parent_scope(node)),
234                            tree.Class,
235                        ):
236                            return []
237                        else:
238                            return None
239                    if (
240                        node.is_definition()
241                        and not global_encountered
242                        and (
243                            is_function_definition
244                            or parser_utils.get_parent_scope(node)
245                            != parser_utils.get_parent_scope(definition)
246                        )
247                    ):
248                        return None
249                    if self._is_name_function_definition(definition) and isinstance(
250                        parser_utils.get_parent_scope(parser_utils.get_parent_scope(definition)),
251                        tree.Class,
252                    ):
253                        return None
254                names.append(node)
255            return names
256
257        if definition:
258            if self._is_name_function_definition(definition):
259                scope = parser_utils.get_parent_scope(parser_utils.get_parent_scope(definition))
260            else:
261                scope = parser_utils.get_parent_scope(definition)
262        else:
263            scope = parser_utils.get_parent_scope(name)
264
265        usages = find_usages_in_node(scope)
266        return usages
267
268    def get_positions_for(self, source, line, column):
269        module_node = jedi_utils.parse_source(source)
270        pos = (line, column)
271        stmt = self._get_statement_for_position(module_node, pos)
272
273        name = None
274        if isinstance(stmt, tree.Name):
275            name = stmt
276        elif isinstance(stmt, tree.BaseNode):
277            name = stmt.get_name_of_position(pos)
278
279        if not name:
280            return set()
281
282        # format usage positions as tkinter text widget indices
283        return set(
284            (
285                "%d.%d" % (usage.start_pos[0], usage.start_pos[1]),
286                "%d.%d" % (usage.start_pos[0], usage.start_pos[1] + len(name.value)),
287            )
288            for usage in self._find_usages(name, stmt)
289        )
290
291
292class UsagesHighlighter(BaseNameHighlighter):
293    """Script.usages looks tempting method to use for finding variable occurrences,
294    but it only returns last
295    assignments to a variable, not really all usages (with Jedi 0.10).
296    But it finds attribute usages quite nicely.
297
298    TODO: check if this gets fixed in later versions of Jedi
299
300    NB!!!!!!!!!!!!! newer jedi versions use subprocess and are too slow to run
301    for each keypress"""
302
303    def get_positions_for(self, source, line, column):
304        # https://github.com/davidhalter/jedi/issues/897
305        from jedi import Script
306
307        script = Script(source + ")")
308        usages = script.get_references(line, column, include_builtins=False)
309
310        result = {
311            (
312                "%d.%d" % (usage.line, usage.column),
313                "%d.%d" % (usage.line, usage.column + len(usage.name)),
314            )
315            for usage in usages
316            if usage.module_name == ""
317        }
318
319        return result
320
321
322class CombinedHighlighter(VariablesHighlighter, UsagesHighlighter):
323    def get_positions_for(self, source, line, column):
324        usages = UsagesHighlighter.get_positions_for(self, source, line, column)
325        variables = VariablesHighlighter.get_positions_for(self, source, line, column)
326        return usages | variables
327
328
329def update_highlighting(event):
330    if not get_workbench().ready:
331        # don't slow down loading process
332        return
333
334    global tree
335    if not tree:
336        # using lazy importing to speed up Thonny loading
337        from parso.python import tree  # pylint: disable=redefined-outer-name
338
339    assert isinstance(event.widget, tk.Text)
340    text = event.widget
341
342    if not hasattr(text, "name_highlighter"):
343        text.name_highlighter = VariablesHighlighter(text)
344        # Alternatives:
345        # NB! usages() is too slow when used on library names
346        # text.name_highlighter = CombinedHighlighter(text)
347        # text.name_highlighter = UsagesHighlighter(text)
348
349    text.name_highlighter.schedule_update()
350
351
352def load_plugin() -> None:
353    wb = get_workbench()
354    wb.set_default("view.name_highlighting", False)
355    wb.bind_class("CodeViewText", "<<CursorMove>>", update_highlighting, True)
356    wb.bind_class("CodeViewText", "<<TextChange>>", update_highlighting, True)
357    wb.bind("<<UpdateAppearance>>", update_highlighting, True)
358