injector.py

import string
import threading
import time
from dataclasses import dataclass, field

from pynput.keyboard import Controller, Key

from corrector import correct

BOUNDARY_KEYS = {Key.space, Key.enter}
WORD_CHARS = set(string.ascii_letters + string.digits + "'")


@dataclass
class AutoCorrectEngine:
    enabled: bool = True
    keyboard: Controller = field(default_factory=Controller)
    current_word: list[str] = field(default_factory=list)
    lock: threading.Lock = field(default_factory=threading.Lock)
    suppress_until: float = 0.0

    def is_enabled(self) -> bool:
        with self.lock:
            return self.enabled

    def set_enabled(self, value: bool) -> bool:
        with self.lock:
            self.enabled = value
            if not value:
                self.current_word.clear()
            return self.enabled

    def toggle(self) -> bool:
        with self.lock:
            self.enabled = not self.enabled
            if not self.enabled:
                self.current_word.clear()
            return self.enabled

    def should_ignore_event(self) -> bool:
        return time.monotonic() < self.suppress_until

    def handle_press(self, key) -> None:
        if self.should_ignore_event():
            return

        if key in BOUNDARY_KEYS:
            self._handle_boundary(key)
            return

        if key == Key.backspace:
            with self.lock:
                if self.current_word:
                    self.current_word.pop()
            return

        if key == Key.esc:
            with self.lock:
                self.current_word.clear()
            return

        if key in {
            Key.delete,
            Key.left,
            Key.right,
            Key.up,
            Key.down,
            Key.home,
            Key.end,
            Key.page_up,
            Key.page_down,
        }:
            with self.lock:
                self.current_word.clear()
            return

        char = getattr(key, "char", None)
        if char is None:
            return

        if char in WORD_CHARS:
            with self.lock:
                self.current_word.append(char)
        else:
            with self.lock:
                self.current_word.clear()

    def _handle_boundary(self, key) -> None:
        with self.lock:
            if not self.enabled:
                self.current_word.clear()
                return
            word = "".join(self.current_word)
            self.current_word.clear()

        if len(word) < 2:
            return

        corrected = correct(word)
        if corrected == word:
            return

        replacement = corrected + self._boundary_text(key)
        self._replace_last_word(word, replacement)
        print(f"[autocorrect] {word} -> {corrected}")

    def _replace_last_word(self, original_word: str, replacement: str) -> None:
        # Ignore key events generated by our own backspaces and retyped text.
        self.suppress_until = time.monotonic() + 0.6

        for _ in range(len(original_word) + 1):
            self.keyboard.press(Key.backspace)
            self.keyboard.release(Key.backspace)
            time.sleep(0.005)

        self.keyboard.type(replacement)

    @staticmethod
    def _boundary_text(key) -> str:
        if key == Key.enter:
            return "\n"
        return " "