better faulty encoding detection

copienator/latex_json.py

@@ -13,6 +13,9 @@ _CONTENT_FIELDS = {'question_content', 'solution_content', 'context_content', 'r
 _CONTENT_KEYS = {f'"{field}"' for field in _CONTENT_FIELDS}
 _STRING = re.compile(r'"(?:\\.|[^"\\])*"', re.DOTALL)
 _COMMAND = re.compile(r'[a-zA-Z]+')
+_MATH = re.compile(r"(\$\$.*?\$\$|\$.*?\$|\\\(.*?\\\)|\\\[.*?\\\])", re.DOTALL)
+_DECODED_COMMAND = re.compile(r"[\b\f\n\r\t]([A-Za-z]+)")
+_JSON_CONTROL_PREFIX = {"\b": "b", "\f": "f", "\n": "n", "\r": "r", "\t": "t"}
 
 
 def normalize_literal_newlines(text: str) -> str:
@@ -44,6 +47,29 @@ underbrace underline underset uparrow updownarrow upsilon usepackage
 '''.split())
 
 
+def repair_decoded_latex_controls(text: str) -> str:
+    r"""Restore known TeX commands after JSON decoded their first letter.
+
+    Only math spans and whitelisted command names are changed. This avoids
+    treating ordinary paragraph newlines or tabs as damaged LaTeX.
+    """
+    def repair_math(math_match: re.Match[str]) -> str:
+        def repair_command(command_match: re.Match[str]) -> str:
+            candidate = (
+                _JSON_CONTROL_PREFIX[command_match.group(0)[0]]
+                + command_match.group(1)
+            )
+            return (
+                f"\\{candidate}"
+                if candidate in _AMBIGUOUS_COMMANDS
+                else command_match.group(0)
+            )
+
+        return _DECODED_COMMAND.sub(repair_command, math_match.group(0))
+
+    return _MATH.sub(repair_math, text)
+
+
 def _repair_string(token: str) -> str:
     output = []
     index = 0
@@ -71,7 +97,9 @@ def _repair_string(token: str) -> str:
     return ''.join(output)
 
 
-def repair_latex_json(raw: str) -> str:
+def repair_latex_json(
+    raw: str, *, additional_content_fields: set[str] | None = None
+) -> str:
     """Repair single TeX backslashes without decoding or rewriting other JSON.
 
     Preserve standard escapes, Unicode escapes, and already escaped TeX.
@@ -79,14 +107,22 @@ def repair_latex_json(raw: str) -> str:
     as a broken ``\\not`` command; known TeX command names take precedence.
     Raw responses/caches should be retained to allow reviewing such cases.
     """
+    additional_content_keys = {
+        f'"{field}"' for field in (additional_content_fields or set())
+    }
     output = []
     end = 0
     previous = None
     for match in _STRING.finditer(raw):
         token = match.group()
-        if (previous is not None and raw[previous.end():match.start()].strip() == ':'
-                and previous.group() in _CONTENT_KEYS):
-            token = _repair_string(token)
+        if previous is not None and raw[previous.end():match.start()].strip() == ':':
+            if previous.group() in _CONTENT_KEYS:
+                token = _repair_string(token)
+            elif previous.group() in additional_content_keys:
+                token = _MATH.sub(
+                    lambda math_match: _repair_string(math_match.group(0)),
+                    token,
+                )
         output.extend((raw[end:match.start()], token))
         end = match.end()
         previous = match