1import contextlib
2import re
3import string
4
5import urwid
6
7
8def _is_valid_key(char):
9    return urwid.util.is_wide_char(char, 0) or (
10        len(char) == 1 and ord(char) >= 32
11    )
12
13
14class AutocompleteState:
15    def __init__(self, prefix, infix, suffix, cycle_forward):
16        self.prefix = prefix
17        self.infix = infix
18        self.suffix = suffix
19        self.num = 0 if cycle_forward else -1
20
21
22class PasteBuffer(list):
23    def append(self, text):
24        if not len(text):
25            return
26        super().append(text)
27
28
29class UndoState:
30    def __init__(self, edit_pos, edit_text):
31        self.edit_pos = edit_pos
32        self.edit_text = edit_text
33
34
35class UndoBuffer:
36    def __init__(self):
37        self.pos = 0
38        self.buffer = []
39
40    @property
41    def empty(self):
42        return self.pos == 0
43
44    @property
45    def cur(self):
46        return self.buffer[self.pos - 1]
47
48    def push(self, old_state, new_state):
49        self.buffer = self.buffer[: self.pos]
50        if old_state.edit_text != new_state.edit_text:
51            self.buffer.append((old_state, new_state))
52        self.pos = len(self.buffer)
53
54    def pop(self):
55        if not self.empty:
56            self.pos -= 1
57
58
59class ReadlineEdit(urwid.Edit):
60    ignore_focus = False
61
62    def __init__(
63        self,
64        *args,
65        word_chars=string.ascii_letters + string.digits + "_",
66        max_char=None,
67        **kwargs
68    ):
69        if max_char and "edit_text" in kwargs:
70            kwargs["edit_text"] = kwargs["edit_text"][:max_char]
71        super().__init__(*args, **kwargs)
72        self._word_regex1 = re.compile(
73            "([%s]+)" % "|".join(re.escape(ch) for ch in word_chars)
74        )
75        self._word_regex2 = re.compile(
76            "([^%s]+)" % "|".join(re.escape(ch) for ch in word_chars)
77        )
78        self._autocomplete_state = None
79        self._autocomplete_func = None
80        self._autocomplete_key = None
81        self._autocomplete_key_reverse = None
82        self._autocomplete_delims = " \t\n;"
83        self._max_char = max_char
84        self._paste_buffer = PasteBuffer()
85        self._undo_buffer = UndoBuffer()
86        self.size = (30,)  # SET MAXCOL DEFAULT VALUE
87
88        self.keymap = {
89            "ctrl f": self.forward_char,
90            "ctrl b": self.backward_char,
91            "ctrl a": self.beginning_of_line,
92            "ctrl e": self.end_of_line,
93            "home": self.beginning_of_line,
94            "end": self.end_of_line,
95            "meta f": self.forward_word,
96            "meta b": self.backward_word,
97            "shift right": self.forward_word,
98            "shift left": self.backward_word,
99            "ctrl d": self.delete_char,
100            "ctrl h": self.backward_delete_char,
101            "delete": self.delete_char,
102            "backspace": self.backward_delete_char,
103            "ctrl u": self.backward_kill_line,
104            "ctrl k": self.forward_kill_line,
105            "meta x": self.kill_whole_line,
106            "meta d": self.kill_word,
107            "ctrl w": self.backward_kill_word,
108            "meta backspace": self.backward_kill_word,
109            "ctrl t": self.transpose_chars,
110            "ctrl l": self.clear_screen,
111            "ctrl y": self.paste,
112            "ctrl _": self.undo,
113        }
114
115        if self.multiline:
116            self.keymap.update(
117                {
118                    "up": self.previous_line,
119                    "ctrl p": self.previous_line,
120                    "ctrl n": self.next_line,
121                    "down": self.next_line,
122                    "enter": self.insert_new_line,
123                }
124            )
125
126    def keypress(self, size, key):
127        self.size = size
128        if key == self._autocomplete_key and self._autocomplete_func:
129            self._complete(True)
130            return None
131        elif key == self._autocomplete_key_reverse and self._autocomplete_func:
132            self._complete(False)
133            return None
134        else:
135            self._autocomplete_state = None
136
137        if key == "right":
138            return None if self.forward_char() else key
139
140        if key == "left":
141            return None if self.backward_char() else key
142
143        if key in self.keymap:
144            if self.keymap[key] == self.undo:
145                self.keymap[key]()
146            else:
147                with self._capture_undo():
148                    self.keymap[key]()
149            self._invalidate()
150            return None
151        elif _is_valid_key(key):
152            with self._capture_undo():
153                self._insert_char_at_cursor(key)
154            self._invalidate()
155            return None
156        return key
157
158    def _insert_char_at_cursor(self, key):
159        if self._max_char and len(self.edit_text) == self._max_char:
160            return
161
162        self.set_edit_text(
163            self._edit_text[0 : self._edit_pos]
164            + key
165            + self._edit_text[self._edit_pos :]
166        )
167        self.set_edit_pos(self._edit_pos + 1)
168
169    def clear_screen(self):
170        self.set_edit_pos(0)
171        self.set_edit_text("")
172
173    def _make_undo_state(self):
174        return UndoState(self.edit_pos, self.edit_text)
175
176    def _apply_undo_state(self, state):
177        self.set_edit_text(state.edit_text)
178        self.set_edit_pos(state.edit_pos)
179
180    @contextlib.contextmanager
181    def _capture_undo(self):
182        old_state = self._make_undo_state()
183        yield
184        new_state = self._make_undo_state()
185        self._undo_buffer.push(old_state, new_state)
186
187    def undo(self):
188        if self._undo_buffer.empty:
189            return
190        old_state, new_state = self._undo_buffer.cur
191        self._undo_buffer.pop()
192        self._apply_undo_state(old_state)
193
194    def paste(self):
195        # do not paste if empty buffer
196        if not len(self._paste_buffer):
197            return
198
199        text = self._paste_buffer[-1]
200        if self._max_char:
201            chars_left = self._max_char - len(self.edit_text)
202            text = text[:chars_left]
203
204        self.set_edit_text(
205            self.edit_text[: self.edit_pos]
206            + text
207            + self.edit_text[self.edit_pos :]
208        )
209        self.set_edit_pos(self.edit_pos + len(text))
210
211    def previous_line(self):
212        x, y = self.get_cursor_coords(self.size)
213        self.move_cursor_to_coords(self.size, x, max(0, y - 1))
214
215    def next_line(self):
216        x, y = self.get_cursor_coords(self.size)
217        self.move_cursor_to_coords(self.size, x, y + 1)
218
219    def backward_char(self):
220        if self._edit_pos > 0:
221            self.set_edit_pos(self._edit_pos - 1)
222            return True
223        return False
224
225    def forward_char(self):
226        if self._edit_pos < len(self._edit_text):
227            self.set_edit_pos(self._edit_pos + 1)
228            return True
229        return False
230
231    def backward_word(self):
232        for match in self._word_regex1.finditer(
233            self._edit_text[0 : self._edit_pos][::-1]
234        ):
235            self.set_edit_pos(self._edit_pos - match.end(1))
236            return
237        self.set_edit_pos(0)
238
239    def forward_word(self):
240        for match in self._word_regex2.finditer(
241            self._edit_text[self._edit_pos :]
242        ):
243            self.set_edit_pos(self._edit_pos + match.end(1))
244            return
245        self.set_edit_pos(len(self._edit_text))
246
247    def delete_char(self):
248        if self._edit_pos < len(self._edit_text):
249            self.set_edit_text(
250                self._edit_text[0 : self._edit_pos]
251                + self._edit_text[self._edit_pos + 1 :]
252            )
253
254    def backward_delete_char(self):
255        if self._edit_pos > 0:
256            self.set_edit_pos(self._edit_pos - 1)
257            self.set_edit_text(
258                self._edit_text[0 : self._edit_pos]
259                + self._edit_text[self._edit_pos + 1 :]
260            )
261
262    def backward_kill_line(self):
263        for pos in reversed(range(0, self.edit_pos)):
264            if self.edit_text[pos] == "\n":
265                self._paste_buffer.append(
266                    self.edit_text[pos + 1 : self.edit_pos]
267                )
268                self.set_edit_text(
269                    self._edit_text[: pos + 1]
270                    + self._edit_text[self.edit_pos :]
271                )
272                self.edit_pos = pos + 1
273                return
274        self._paste_buffer.append(self.edit_text[: self.edit_pos])
275        self.set_edit_text(self._edit_text[self.edit_pos :])
276        self.edit_pos = 0
277
278    def forward_kill_line(self):
279        for pos in range(self.edit_pos, len(self.edit_text)):
280            if self.edit_text[pos] == "\n":
281                self._paste_buffer.append(self.edit_text[self.edit_pos : pos])
282                self.set_edit_text(
283                    self._edit_text[: self.edit_pos] + self._edit_text[pos:]
284                )
285                return
286        self._paste_buffer.append(self.edit_text[self.edit_pos :])
287        self.set_edit_text(self._edit_text[: self.edit_pos])
288
289    def kill_whole_line(self):
290        buffer_length = len(self._paste_buffer)
291        self.backward_kill_line()
292        self.forward_kill_line()
293        if len(self._paste_buffer) - buffer_length == 2:
294            # if text was added from both forward and backward kill
295            self._paste_buffer[:2] = ["".join(self._paste_buffer[:2])]
296
297    def backward_kill_word(self):
298        pos = self._edit_pos
299        self.backward_word()
300        self._paste_buffer.append(self._edit_text[self.edit_pos : pos])
301        self.set_edit_text(
302            self._edit_text[: self._edit_pos] + self._edit_text[pos:]
303        )
304
305    def kill_word(self):
306        pos = self._edit_pos
307        self.forward_word()
308        self._paste_buffer.append(self.edit_text[pos : self.edit_pos])
309        self.set_edit_text(
310            self._edit_text[0:pos] + self._edit_text[self._edit_pos :]
311        )
312        self.set_edit_pos(pos)
313
314    def beginning_of_line(self):
315        x, y = self.get_cursor_coords(self.size)
316        if x == 0 and y > 0:
317            y -= 1
318        self.move_cursor_to_coords(self.size, 0, y)
319
320    def end_of_line(self):
321        text_length = len(self.edit_text)
322        # Move one character forward if at the end of a line.
323        if (
324            self.edit_pos < text_length
325            and self.edit_text[self.edit_pos] == "\n"
326        ):
327            self.forward_char()
328        # Set the position of cursor at the next '\n'.
329        for pos in range(self.edit_pos, text_length + 1):
330            if pos == text_length:
331                self.set_edit_pos(pos)
332                return
333            elif self.edit_text[pos] == "\n":
334                self.set_edit_pos(pos)
335                return
336
337    def transpose_chars(self):
338        x, y = self.get_cursor_coords(self.size)
339        x = max(2, x + 1)
340        self.move_cursor_to_coords(self.size, x, y)
341        x, y = self.get_cursor_coords(self.size)
342        if x == 1:
343            # Don't transpose in case of single character
344            return
345        self.set_edit_text(
346            self._edit_text[0 : self._edit_pos - 2]
347            + self._edit_text[self._edit_pos - 1]
348            + self._edit_text[self._edit_pos - 2]
349            + self._edit_text[self._edit_pos :]
350        )
351
352    def insert_new_line(self):
353        if self.multiline:
354            self.insert_text("\n")
355
356    def enable_autocomplete(self, func, key="tab", key_reverse="shift tab"):
357        self._autocomplete_func = func
358        self._autocomplete_key = key
359        self._autocomplete_key_reverse = key_reverse
360        self._autocomplete_state = None
361
362    def set_completer_delims(self, delimiters):
363        self._autocomplete_delims = delimiters
364
365    def _complete(self, cycle_forward):
366        if self._autocomplete_state:
367            if self._autocomplete_state.num == 0 and not cycle_forward:
368                self._autocomplete_state.num = None
369            elif self._autocomplete_state.num == -1 and cycle_forward:
370                self._autocomplete_state.num = None
371            else:
372                self._autocomplete_state.num += 1 if cycle_forward else -1
373        else:
374            text_before_caret = self.edit_text[0 : self.edit_pos]
375            text_after_caret = self.edit_text[self.edit_pos :]
376
377            if self._autocomplete_delims:
378                group = re.escape(self._autocomplete_delims)
379                match = re.match(
380                    "^(?P<prefix>.*[" + group + "])(?P<infix>.*?)$",
381                    text_before_caret,
382                    flags=re.M | re.DOTALL,
383                )
384                if match:
385                    prefix = match.group("prefix")
386                    infix = match.group("infix")
387                else:
388                    prefix = ""
389                    infix = text_before_caret
390            else:
391                match = re.match(
392                    "^(?P<infix>.*?)$",
393                    text_before_caret,
394                    flags=re.M | re.DOTALL,
395                )
396                prefix = ""
397                if match:
398                    infix = match.group("infix")
399                else:
400                    infix = text_before_caret
401
402            suffix = text_after_caret
403
404            self._autocomplete_state = AutocompleteState(
405                prefix, infix, suffix, cycle_forward
406            )
407
408        state = self._autocomplete_state
409
410        match = self._autocomplete_func(state.infix, state.num)
411        if not match:
412            match = state.infix
413            self._autocomplete_state = None
414
415        self.edit_text = state.prefix + match + state.suffix
416        self.edit_pos = len(state.prefix) + len(match)
417