fix: preserve custom LaTeX env shorthand macros during translation

crazy_functions/latex_fns/latex_actions.py

@@ -9,7 +9,7 @@
 from crazy_functions.latex_fns.latex_toolbox import set_forbidden_text, set_forbidden_text_begin_end, set_forbidden_text_careful_brace
 from crazy_functions.latex_fns.latex_toolbox import reverse_forbidden_text_careful_brace, reverse_forbidden_text, convert_to_linklist, post_process
 from crazy_functions.latex_fns.latex_toolbox import fix_content, find_main_tex_file, merge_tex_files, compile_latex_with_timeout
-from crazy_functions.latex_fns.latex_toolbox import find_title_and_abs
+from crazy_functions.latex_fns.latex_toolbox import find_title_and_abs, extract_user_defined_env_patterns
 from crazy_functions.latex_fns.latex_pickle_io import objdump, objload
 
 
@@ -48,6 +48,10 @@ def split_subprocess(txt, project_folder, return_dict, opts):
     text, mask = set_forbidden_text(text, mask, [r"\\begin\{minipage\}(.*?)\\end\{minipage\}", r"\\begin\{minipage\*\}(.*?)\\end\{minipage\*\}"], re.DOTALL)
     text, mask = set_forbidden_text(text, mask, [r"\\begin\{align\*\}(.*?)\\end\{align\*\}", r"\\begin\{align\}(.*?)\\end\{align\}"], re.DOTALL)
     text, mask = set_forbidden_text(text, mask, [r"\\begin\{equation\}(.*?)\\end\{equation\}", r"\\begin\{equation\*\}(.*?)\\end\{equation\*\}"], re.DOTALL)
+    # 吸收用户自定义的环境简写命令 (如 \def\be{\begin{equation}} \def\ee{\end{equation}})
+    custom_env_patterns = extract_user_defined_env_patterns(text)
+    if custom_env_patterns:
+        text, mask = set_forbidden_text(text, mask, custom_env_patterns, re.DOTALL)
     text, mask = set_forbidden_text(text, mask, [r"\\includepdf\[(.*?)\]\{(.*?)\}", r"\\clearpage", r"\\newpage", r"\\appendix", r"\\tableofcontents", r"\\include\{(.*?)\}"])
     text, mask = set_forbidden_text(text, mask, [r"\\vspace\{(.*?)\}", r"\\hspace\{(.*?)\}", r"\\label\{(.*?)\}", r"\\begin\{(.*?)\}", r"\\end\{(.*?)\}", r"\\item "])
     text, mask = set_forbidden_text_careful_brace(text, mask, r"\\hl\{(.*?)\}", re.DOTALL)

crazy_functions/latex_fns/latex_toolbox.py

@@ -592,6 +592,41 @@ def find_next(string, chars, begin):
     return final_tex
 
 
+def extract_user_defined_env_patterns(text):
+    """
+    Detect user-defined shorthand commands for LaTeX environments, e.g.:
+        \\def\\be{\\begin{equation}}
+        \\def\\ee{\\end{equation}}
+    Returns a list of regex patterns that can be passed to set_forbidden_text
+    to preserve the corresponding environment blocks from GPT translation.
+    """
+    begin_defs = {}  # shorthand -> env_name
+    end_defs = {}    # shorthand -> env_name
+
+    # Match \\def\\cmd{\\begin{envname}}
+    for m in re.finditer(r'\\def\\([a-zA-Z]+)\s*\{\\begin\{([^}]+)\}\}', text):
+        begin_defs[m.group(1)] = m.group(2)
+    # Match \\def\\cmd{\\end{envname}}
+    for m in re.finditer(r'\\def\\([a-zA-Z]+)\s*\{\\end\{([^}]+)\}\}', text):
+        end_defs[m.group(1)] = m.group(2)
+
+    # Match \\newcommand{\\cmd}{\\begin{envname}} and \\renewcommand variants
+    for m in re.finditer(r'\\(?:newcommand|renewcommand)\{\\([a-zA-Z]+)\}\{\\begin\{([^}]+)\}\}', text):
+        begin_defs[m.group(1)] = m.group(2)
+    for m in re.finditer(r'\\(?:newcommand|renewcommand)\{\\([a-zA-Z]+)\}\{\\end\{([^}]+)\}\}', text):
+        end_defs[m.group(1)] = m.group(2)
+
+    patterns = []
+    for begin_cmd, env_name in begin_defs.items():
+        for end_cmd, end_env_name in end_defs.items():
+            if env_name == end_env_name:
+                # Use (?![a-zA-Z]) to avoid matching longer commands that start with the same prefix
+                # e.g., \be should not match \begin
+                patterns.append(rf'\\{begin_cmd}(?![a-zA-Z])(.*?)\\{end_cmd}(?![a-zA-Z])')
+                logger.info(f"Detected custom env shorthand: \\{begin_cmd}...\\{end_cmd} -> {{{env_name}}}")
+    return patterns
+
+
 def compile_latex_with_timeout(command, cwd, timeout=60):
     import subprocess