sympype: pipe LaTeX math expressions through Sympy right from within VSCode

sympype.py

#!/usr/bin/env python

# sympype.py 0.0.2 (2022-03-19)
###############################
# This script is intended to be used as a "pipe filter" in
# Vim or VSCode with Vim Extension. Suppose the LaTeX file
# I'm editing has an algebra exercise consisting of expressions
# that are to be simplified:
#
# \item $\frac{x^2-9}{x+3}$
# \item $\frac{x^2+x}{x}$
#
# The expressions can be replaced by the expected solutions
# by piping the whole block through sympype:
# - visually select the lines (Shift+V in command mode)
# - enter ":! sympype.py 'simplify(X)'" as Vim command
# Presto! The example expressions have been replaced by:
#
# \item $x - 3$
# \item $x + 1$
#
# NOTE:
# This script has no error checking and has been written
# exclusively for (my) personal use.
# I do not intend to develop it further or do active maintenance.
# TODO: support for display math

import sys
import re
from sympy import latex
from sympy.parsing.sympy_parser import parse_expr
from sympy.parsing.latex import parse_latex


def strip_latex_comments(s):
    return re.sub(r'%.*', '', s)


def cleanup_latex_math(s: str):
    regexes = [
        (r'\\ds\s?', ''),              # remove custom '\ds' macro
        (r'[^a-z]([a-z])\(', r'\1*('), # "fix" implicit multiplication
        (r'\\\(', r'\\left('),         # \( is my custom macro for \left(
        (r'\\\)', r'\\right)'),
    ]
    for pat, repl in regexes:
        s = re.sub(pat, repl, s)
    return s


def process_mathmode(matchobj, sympymanip: str):
    """replace a single LaTeX math expression by its evaluated version"""
    s = cleanup_latex_math(matchobj.group(1))
    sympyex = parse_latex(s)
    sympyres = parse_expr(sympymanip, local_dict={'X': sympyex})
    res_latex = latex(sympyres)
    return '$' + res_latex + '$'


def apply_sympy_to_latex(multilinelatex: str, sympymanip: str):
    """ Apply given Sympy manipulation to all inline math in LaTeX string

        multilinelatex: complete LaTeX code block
        sympymanip: sympy code containing 'X' as placeholder"""
    s = strip_latex_comments(multilinelatex)
    def proc_math(m): return process_mathmode(m, sympymanip)
    # mathmode may be split over more than one line, use re.DOTALL:
    out = re.sub(r'\$(.*?)\$', proc_math, s, flags=re.DOTALL)
    return out


if __name__ == "__main__":
    sympymanip = sys.argv[1]
    multilinelatex = sys.stdin.read()
    print(apply_sympy_to_latex(multilinelatex, sympymanip))