ocr_studio/language.py

from __future__ import annotations

import re

from ocr_studio.config import LANGUAGE_AUTO, LANGUAGE_BOTH, LANGUAGE_EN, LANGUAGE_FA

ARABIC_SCRIPT_RE = re.compile(r"[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF]")
LATIN_RE = re.compile(r"[A-Za-z]")

_ARABIC_YE = "\u064a"
_PERSIAN_YE = "\u06cc"
_ARABIC_KEHEH = "\u0643"
_PERSIAN_KEHEH = "\u06a9"


def contains_arabic_script(text: str) -> bool:
    return bool(ARABIC_SCRIPT_RE.search(text or ""))


def detect_language(text: str) -> str:
    arabic = len(ARABIC_SCRIPT_RE.findall(text or ""))
    latin = len(LATIN_RE.findall(text or ""))
    if arabic and latin:
        return LANGUAGE_BOTH
    if arabic:
        return LANGUAGE_FA
    if latin:
        return LANGUAGE_EN
    return LANGUAGE_BOTH


def resolve_language(selected: str, text: str) -> str:
    if selected == LANGUAGE_AUTO:
        return detect_language(text)
    if selected in {LANGUAGE_BOTH, LANGUAGE_FA, LANGUAGE_EN}:
        return selected
    return LANGUAGE_BOTH


def normalize_persian(text: str) -> str:
    return (
        (text or "")
        .replace(_ARABIC_YE, _PERSIAN_YE)
        .replace(_ARABIC_KEHEH, _PERSIAN_KEHEH)
    )


def normalize_text(text: str, language: str) -> str:
    cleaned = (text or "").replace("\r\n", "\n").replace("\r", "\n")
    cleaned = re.sub(r"[ \t]+\n", "\n", cleaned)
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip()
    if language in {LANGUAGE_FA, LANGUAGE_BOTH, LANGUAGE_AUTO}:
        cleaned = normalize_persian(cleaned)
    return cleaned


def is_rtl_language(language: str) -> bool:
    return language in {LANGUAGE_FA, LANGUAGE_BOTH}