1import logging
2import tkinter as tk
3
4from thonny import get_workbench, jedi_utils
5
6
7class LocalsHighlighter:
8    def __init__(self, text):
9        self.text = text
10
11        self._update_scheduled = False
12
13    def get_positions(self):
14        from jedi import parser_utils
15        from parso.python import tree
16
17        locs = []
18
19        def process_scope(scope):
20            if isinstance(scope, tree.Function):
21                # process all children after name node,
22                # (otherwise name of global function will be marked as local def)
23                local_names = set()
24                global_names = set()
25                for child in scope.children[2:]:
26                    process_node(child, local_names, global_names)
27            else:
28                if hasattr(scope, "subscopes"):
29                    for child in scope.subscopes:
30                        process_scope(child)
31                elif hasattr(scope, "children"):
32                    for child in scope.children:
33                        process_scope(child)
34
35        def process_node(node, local_names, global_names):
36            if isinstance(node, tree.GlobalStmt):
37                global_names.update([n.value for n in node.get_global_names()])
38
39            elif isinstance(node, tree.Name):
40                if node.value in global_names:
41                    return
42
43                if node.is_definition():  # local def
44                    locs.append(node)
45                    local_names.add(node.value)
46                elif node.value in local_names:  # use of local
47                    locs.append(node)
48
49            elif isinstance(node, tree.BaseNode):
50                # ref: jedi/parser/grammar*.txt
51                if node.type == "trailer" and node.children[0].value == ".":
52                    # this is attribute
53                    return
54
55                if isinstance(node, tree.Function):
56                    global_names = set()  # outer global statement doesn't have effect anymore
57
58                for child in node.children:
59                    process_node(child, local_names, global_names)
60
61        source = self.text.get("1.0", "end")
62        module = jedi_utils.parse_source(source)
63        for child in module.children:
64            if isinstance(child, tree.BaseNode) and parser_utils.is_scope(child):
65                process_scope(child)
66
67        loc_pos = set(
68            (
69                "%d.%d" % (usage.start_pos[0], usage.start_pos[1]),
70                "%d.%d" % (usage.start_pos[0], usage.start_pos[1] + len(usage.value)),
71            )
72            for usage in locs
73        )
74
75        return loc_pos
76
77    def _highlight(self, pos_info):
78        for pos in pos_info:
79            start_index, end_index = pos[0], pos[1]
80            self.text.tag_add("local_name", start_index, end_index)
81
82    def schedule_update(self):
83        def perform_update():
84            try:
85                self.update()
86            finally:
87                self._update_scheduled = False
88
89        if not self._update_scheduled:
90            self._update_scheduled = True
91            self.text.after_idle(perform_update)
92
93    def update(self):
94        self.text.tag_remove("local_name", "1.0", "end")
95
96        if get_workbench().get_option("view.locals_highlighting") and self.text.is_python_text():
97            try:
98                highlight_positions = self.get_positions()
99                self._highlight(highlight_positions)
100            except Exception:
101                logging.exception("Problem when updating local variable tags")
102
103
104def update_highlighting(event):
105    if not get_workbench().ready:
106        # don't slow down loading process
107        return
108
109    assert isinstance(event.widget, tk.Text)
110    text = event.widget
111
112    if not hasattr(text, "local_highlighter"):
113        text.local_highlighter = LocalsHighlighter(text)
114
115    text.local_highlighter.schedule_update()
116
117
118def load_plugin() -> None:
119    wb = get_workbench()
120    wb.set_default("view.locals_highlighting", False)
121    wb.bind_class("CodeViewText", "<<TextChange>>", update_highlighting, True)
122    wb.bind("<<UpdateAppearance>>", update_highlighting, True)
123