corrector.py

from pathlib import Path
import sys

from symspellpy import SymSpell, Verbosity
import symspellpy

DEFAULT_OVERRIDES = {
    "teh": "the",
    "adn": "and",
    "nad": "and",
    "wnat": "want",
    "waht": "what",
    "dont": "don't",
    "cant": "can't",
    "wont": "won't",
    "doesnt": "doesn't",
    "isnt": "isn't",
    "arent": "aren't",
    "didnt": "didn't",
    "somthing": "something",
    "smething": "something",
    "thsi": "this",
    "recieve": "receive",
}

_runtime_ignored_words: set[str] = set()
_runtime_overrides: dict[str, str] = {}
MIN_AUTOCORRECT_LENGTH = 4


def _resolve_dictionary_path(filename: str) -> Path:
    candidates = []

    if getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS"):
        candidates.append(Path(sys._MEIPASS) / "symspellpy" / filename)

    package_dir = Path(symspellpy.__file__).resolve().parent
    candidates.append(package_dir / filename)

    for candidate in candidates:
        if candidate.exists():
            return candidate

    searched = ", ".join(str(path) for path in candidates)
    raise FileNotFoundError(
        f"Could not find SymSpell dictionary. Searched: {searched}"
    )


sym_spell = SymSpell(max_dictionary_edit_distance=2, prefix_length=7)
sym_spell.load_dictionary(
    str(_resolve_dictionary_path("frequency_dictionary_en_82_765.txt")),
    term_index=0,
    count_index=1,
)
sym_spell.load_bigram_dictionary(
    str(_resolve_dictionary_path("frequency_bigramdictionary_en_243_342.txt")),
    term_index=0,
    count_index=2,
)


def update_runtime_settings(
    ignored_words: list[str] | None = None,
    custom_replacements: dict[str, str] | None = None,
) -> None:
    global _runtime_ignored_words, _runtime_overrides
    _runtime_ignored_words = {word.strip().lower() for word in ignored_words or [] if word}
    _runtime_overrides = {
        wrong.strip().lower(): corrected.strip()
        for wrong, corrected in (custom_replacements or {}).items()
        if wrong and corrected
    }


def _apply_casing(original: str, corrected: str) -> str:
    if original.isupper():
        return corrected.upper()
    if original[:1].isupper() and original[1:].islower():
        return corrected.capitalize()
    return corrected


def _max_distance_for(word: str) -> int:
    if len(word) <= 4:
        return 1
    return 2


def _should_skip_correction(original: str, suggestion_term: str, distance: int) -> bool:
    if not suggestion_term:
        return True

    if " " in suggestion_term:
        return True

    if len(original) <= 2:
        return True

    if distance > _max_distance_for(original):
        return True

    if len(original) >= 6 and distance == 2:
        shared_prefix = _shared_prefix_length(original.lower(), suggestion_term.lower())
        if shared_prefix < 2:
            return True

    return False


def _shared_prefix_length(left: str, right: str) -> int:
    count = 0
    for left_char, right_char in zip(left, right):
        if left_char != right_char:
            break
        count += 1
    return count


def _lookup_override(word: str) -> str | None:
    lowered = word.lower()
    if lowered in _runtime_ignored_words:
        return word

    if lowered in _runtime_overrides:
        return _apply_casing(word, _runtime_overrides[lowered])

    if lowered in DEFAULT_OVERRIDES:
        return _apply_casing(word, DEFAULT_OVERRIDES[lowered])

    return None


def correct_word(word: str) -> str:
    if not word or not word.strip():
        return word

    override = _lookup_override(word)
    if override is not None:
        return override

    # Short words create too many false positives in system-wide autocorrect.
    if len(word) < MIN_AUTOCORRECT_LENGTH:
        return word

    suggestions = sym_spell.lookup(
        word,
        Verbosity.TOP,
        max_edit_distance=_max_distance_for(word),
        include_unknown=True,
        transfer_casing=True,
    )

    if not suggestions:
        return word

    best = suggestions[0]
    if best.term == word:
        return word

    if _should_skip_correction(word, best.term, best.distance):
        return word

    return best.term


def correct(text: str) -> str:
    if not text or not text.strip():
        return text

    lines = text.split("\n")
    corrected_lines = []

    for line in lines:
        if not line.strip():
            corrected_lines.append(line)
            continue

        words = line.split(" ")
        corrected_words = [correct_word(word) if word else word for word in words]
        corrected_lines.append(" ".join(corrected_words))

    return "\n".join(corrected_lines)