expression-parsing: big-operator task symbols wrap their argument

expression_parsing/CHANGELOG.md

@@ -68,6 +68,7 @@ The first public release. Further tidy-up after the 0.2.0 extraction: previewing
 - `{f(x,y), 2}` is two answers, `f(x,y)` and `2`: `create_expression_set` split at every comma, including those inside brackets.
 - `is_number("2*10^3")` is true: the `^` in `NUMBER_REGEX` was an unescaped anchor, so the `*10^` exponent form never matched.
 - `compute_relative_tolerance_from_significant_decimals` counts only the mantissa for `1.5E3` and `1.5*10^3` (it counted the exponent's digits as significant figures, giving 5e-4 instead of 0.05).
+- (since 0.3.0rc3) A task symbol whose LaTeX starts with a big operator (`\sum`, `\int`, `\prod`, `\lim`, `\bigcup`, `\bigcap`; `BIG_OPERATOR_PREFIXES` in `sympy_parsing`) wraps its argument when applied: with `infsum` as `\sum_{n=1}^{\infty}`, `infsum((-1)^n/n^2)` is an undefined function `infsum` of the summand, which `sympy_to_latex` prints as the operator followed by it. It was the product `infsum*((-1)^n/n^2)`, so the operator was printed inside the fraction, or followed by `\cdot`. Not applied (`2*infsum`), it is still the task's symbol. From compareExpressions#261.
 - `parse_latex` recognises `\pm`/`\mp` without a following space (`1\pm\sqrt{2}` failed to parse), and no longer runs symbol substitution over the LaTeX when expanding them (`\lambda` became `\lamda`).
 
 ## 0.2.0 (internal, not published)

expression_parsing/src/lambdafeedback/expression_parsing/sympy_parsing.py

@@ -15,6 +15,7 @@
 
 import sympy
 from sympy import Basic, Chi, E, Equality, Function, I, Lambda, Set, Sum, Symbol, Tuple, beta, gamma, summation, zeta
+from sympy.core.function import UndefinedFunction
 from sympy.parsing.sympy_parser import (
     _token_splittable,
     auto_number,
@@ -95,6 +96,62 @@ def sympy_symbols(symbols: Iterable[str]) -> dict[str, Symbol]:
     return {name: Symbol(name) for name in symbols}
 
 
+BIG_OPERATOR_PREFIXES = (r"\sum", r"\int", r"\prod", r"\lim", r"\bigcup", r"\bigcap")
+"""LaTeX commands that make a task symbol a prefix operator (see :func:`prefix_operator`)."""
+
+
+def _print_prefix_operator(self: Any, printer: Any) -> str:
+    # Read from a class attribute, so the function pickles (a closure over the LaTeX wouldn't).
+    return f"{self._latex_prefix} {printer.doprint(self.args[0])}"
+
+
+def prefix_operator(name: str, latex: str) -> UndefinedFunction:
+    """An undefined function for a task symbol whose LaTeX is a big operator (``\\sum_{n=1}^{\\infty}``).
+
+    As a ``Symbol``, ``infsum(x)`` would be read as ``infsum*x``, and the
+    operator printed inside fractions or followed by ``\\cdot``. As a function
+    it wraps its argument, and prints as ``latex`` followed by it.
+    """
+    return UndefinedFunction(name, _latex_prefix=latex, _latex=_print_prefix_operator)
+
+
+_OPERATOR_SUFFIX = "__operator"
+
+
+def _prefix_operators(params: ExpressionParams) -> dict[str, UndefinedFunction]:
+    """``{code__operator: prefix_operator(code, latex)}`` for each task symbol whose LaTeX starts with a big operator.
+
+    The code itself stays a ``Symbol``: :func:`_applied_prefix_operators`
+    swaps in the operator only where it is applied.
+    """
+    from .latex import extract_latex  # imported here: latex imports this module
+
+    operators = {}
+    for code, spec in params.symbols.items():
+        if spec.latex is None:
+            continue
+        latex = extract_latex(spec.latex).strip()
+        if latex.startswith(BIG_OPERATOR_PREFIXES):
+            operators[f"{code}{_OPERATOR_SUFFIX}"] = prefix_operator(code, latex)
+    return operators
+
+
+def _applied_prefix_operators(tokens: list[Any], local_dict: dict[str, Any], global_dict: Any) -> list[Any]:
+    """Read a big-operator task symbol followed by ``(`` as its prefix operator (``infsum(x)``).
+
+    Not applied to anything (``2*infsum``), it is still the task's symbol.
+    """
+    result: list[Any] = []
+    values = [value for _, value in tokens]
+    for index, (kind, value) in enumerate(tokens):
+        operator = f"{value}{_OPERATOR_SUFFIX}"
+        if kind == NAME and operator in local_dict and _following(values, index) == "(":
+            result.append((NAME, operator))
+        else:
+            result.append((kind, value))
+    return result
+
+
 UNEVALUATED_SUM = Function("sum")
 """A sum written without an index or limits (``sum(x^n)``): an undefined function of its summand.
 
@@ -168,7 +225,7 @@ def __post_init__(self) -> None:
         params = self.params
         unsplittable = [*self.extra_unsplittable_symbols, *params.reserved_keywords]
         unsplittable += [code for code in params.symbols if len(code) > 1]
-        symbol_dict = {**_base_symbol_dict(params), **sympy_symbols(unsplittable)}
+        symbol_dict = {**_base_symbol_dict(params), **sympy_symbols(unsplittable), **_prefix_operators(params)}
         constants: set[str] = set()
         for symbol, assumption in (*self.extra_symbol_assumptions, *params.symbol_assumptions):
             # Assumptions become Symbol keyword arguments; SymPy stores any identifier.
@@ -235,7 +292,14 @@ def transformations(self) -> tuple[Any, ...]:
         """
         # _keywords_as_names goes after auto_symbol (which leaves keywords alone, and
         # would wrap our Symbol('in') again), before the symbol splitting.
-        leading = (lambda_notation, auto_symbol, _keywords_as_names, repeated_decimals, auto_number)
+        leading = (
+            lambda_notation,
+            auto_symbol,
+            _keywords_as_names,
+            _applied_prefix_operators,
+            repeated_decimals,
+            auto_number,
+        )
         transformations: tuple[Any, ...]
         if self.params.strict_syntax:
             transformations = (*leading, function_exponentiation)

expression_parsing/tests/test_regressions.py

@@ -903,6 +903,61 @@ def test_integrate_with_elementary_functions(self, expr, expected):
         assert self.parse(expr, elementary_functions=True) == expected
 
 
+_BIG_OPERATOR_SYMBOLS = {"infsum": {"latex": r"\sum_{n=1}^{\infty}"}, "pi": {"latex": r"\pi"}}
+
+
+class TestBigOperatorSymbols:
+    """A task symbol with big-operator LaTeX (``infsum`` as ``\\sum``) was a Symbol, so ``infsum(x)`` was ``infsum*x``."""
+
+    def parse(self, response):
+        params = ExpressionParams.from_dict({"elementary_functions": True, "symbols": _BIG_OPERATOR_SYMBOLS})
+        expr = parse_expression(preprocess_expression(response, params), SympyParsingConfig.from_params(params))
+        return expr, params
+
+    def test_symbol_is_a_function_of_its_argument(self):
+        expr, _ = self.parse("infsum((-1)^n / n^2)")
+        assert expr.func.__name__ == "infsum"
+
+    @pytest.mark.parametrize(
+        ("response", "latex"),
+        [
+            ("infsum((-1)^n / n^2)", r"\sum_{n=1}^{\infty} \frac{\left(-1\right)^{n}}{n^{2}}"),
+            (
+                "infsum((4/n^2)*(-1)^n*cos(nx))",
+                r"\sum_{n=1}^{\infty} \frac{4 \cdot \left(-1\right)^{n} \cdot \cos{\left(n \cdot x \right)}}{n^{2}}",
+            ),
+            (
+                "((pi^2)/(3))+infsum((4/n^2)*(-1)^n*cos(nx))",
+                r"\frac{\pi^{2}}{3} + \sum_{n=1}^{\infty} \frac{4 \cdot \left(-1\right)^{n} \cdot \cos{\left(n \cdot x \right)}}{n^{2}}",
+            ),
+        ],
+    )
+    def test_prints_the_operator_before_its_argument(self, response, latex):
+        expr, params = self.parse(response)
+        assert sympy_to_latex(expr, params.symbols, {"mul_symbol": "dot"}) == latex
+
+    @pytest.mark.parametrize("latex", [r"\int_0^1", r"\prod_{k=1}^{n}", r"\lim_{x \to 0}", r"$\bigcup_i$"])
+    def test_every_big_operator_is_a_prefix(self, latex):
+        params = ExpressionParams.from_dict({"symbols": {"op": {"latex": latex}}})
+        assert parse_expression("op(x)", SympyParsingConfig.from_params(params)).func.__name__ == "op"
+
+    def test_not_applied_it_is_still_the_symbol(self):
+        expr, _ = self.parse("2*infsum + infsum(x)")
+        assert expr.has(Symbol("infsum"))
+        assert any(term.func.__name__ == "infsum" for term in expr.args)
+
+    def test_other_symbols_stay_symbols(self):
+        params = ExpressionParams.from_dict({"symbols": {"ab": {"latex": r"\alpha_b"}}})
+        assert parse_expression("ab", SympyParsingConfig.from_params(params)) == Symbol("ab")
+
+    def test_config_round_trips_through_pickle(self):
+        import pickle
+
+        params = ExpressionParams.from_dict({"symbols": _BIG_OPERATOR_SYMBOLS})
+        config = pickle.loads(pickle.dumps(SympyParsingConfig.from_params(params)))
+        assert sympy_to_latex(parse_expression("infsum(x)", config), params.symbols) == r"\sum_{n=1}^{\infty} x"
+
+
 def test_parsing_logs_its_decisions_at_debug_only(caplog):
     import logging