Detect Word equation math segments

src/ai_typewriter/math_format.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+from dataclasses import dataclass
 
 GREEK_AND_SYMBOLS = {
     r"\alpha": "α",
@@ -94,6 +95,23 @@ SUBSCRIPT = str.maketrans({
 })
 
 _GROUP_PATTERN = re.compile(r"([_^])\(([^()]+)\)|([_^])\{([^{}]+)\}|([_^])([A-Za-z0-9+\-=])")
+_MATH_START_PATTERN = re.compile(
+    r"(\\[A-Za-z]+|[A-Za-zΑ-Ωα-ω]+(?:[_^](?:\([^()]+\)|\{[^{}]+\}|[A-Za-z0-9+\-=]))+|[∫∑∞≤≥≠≈α-ωΑ-Ω]|=)"
+)
+_MATH_TOKEN_PATTERN = re.compile(
+    r"\s*(?:"
+    r"\\[A-Za-z]+(?:\([^()]+\)|\{[^{}]+\})?"
+    r"|[A-Za-zΑ-Ωα-ω0-9]+(?:[_^](?:\([^()]+\)|\{[^{}]+\}|[A-Za-z0-9+\-=]))*"
+    r"|[=+\-*/×·≤≥≠≈<>]"
+    r"|[(){}]"
+    r")"
+)
+
+
+@dataclass(frozen=True)
+class MathSegment:
+    kind: str
+    value: str
 
 
 def _translate_script(value: str, marker: str) -> str:
@@ -118,7 +136,7 @@ def format_math_text(text: str, mode: str = "plain") -> str:
     if mode == "plain":
         return text
     if mode not in {"unicode", "unicode_math"}:
-        raise ValueError("math_text_format doit être 'plain' ou 'unicode'")
+        raise ValueError("math_text_format doit être 'plain', 'unicode' ou 'word_equation'")
 
     formatted = text
     # Replace longer commands first so \subseteq wins before \subset.
@@ -132,3 +150,56 @@ def format_math_text(text: str, mode: str = "plain") -> str:
         previous = formatted
         formatted = _GROUP_PATTERN.sub(_replace_script, formatted)
     return formatted
+
+
+def _consume_math_sequence(text: str, start: int) -> int:
+    end = start
+    while end < len(text):
+        if text[end] in "\n.,;:!?":
+            break
+        match = _MATH_TOKEN_PATTERN.match(text, end)
+        if match is None:
+            break
+        end = match.end()
+    return end
+
+
+def split_word_equation_segments(text: str) -> list[MathSegment]:
+    r"""Split text into normal text and Word UnicodeMath equation insertions.
+
+    Equation segments keep Word's linear UnicodeMath syntax (z_1, x^2, \alpha),
+    because Word converts that syntax inside the Alt+= equation editor.
+    """
+    segments: list[MathSegment] = []
+    cursor = 0
+    while cursor < len(text):
+        match = _MATH_START_PATTERN.search(text, cursor)
+        if match is None:
+            if cursor < len(text):
+                segments.append(MathSegment("text", text[cursor:]))
+            break
+
+        start = match.start()
+        if start > cursor:
+            segments.append(MathSegment("text", text[cursor:start]))
+
+        end = _consume_math_sequence(text, start)
+        expression = text[start:end].strip()
+        leading = text[start : start + len(text[start:end]) - len(text[start:end].lstrip())]
+        trailing = text[start:end][len(text[start:end].rstrip()):]
+        if leading:
+            segments.append(MathSegment("text", leading))
+        if expression:
+            segments.append(MathSegment("equation", expression))
+        if trailing:
+            segments.append(MathSegment("text", trailing))
+        cursor = end
+
+    # Merge adjacent text segments for cleaner action counts.
+    merged: list[MathSegment] = []
+    for segment in segments:
+        if merged and segment.kind == "text" and merged[-1].kind == "text":
+            merged[-1] = MathSegment("text", merged[-1].value + segment.value)
+        else:
+            merged.append(segment)
+    return merged