fix: sub operate cause bracket error (#80)

Author: chunibyo-wly

src/integration_tests/regression_test.py

@@ -216,3 +216,11 @@ def sigmoid(x):
         ),
         reduce_assignments=True,
     )
+
+
+def test_sub_bracket() -> None:
+    def solve(a, b):
+        return ((a + b) - b) / (a - b) - (a + b) - (a - b) - (a * b)
+
+    latex = r"\mathrm{solve}(a, b) = \frac{a + b - b}{a - b} - (a + b) - (a - b) - ab"
+    _check_function(solve, latex)

src/latexify/codegen/function_codegen.py

@@ -229,20 +229,20 @@ def visit_BinOp(self, node: ast.BinOp) -> str:
         def _unwrap(child):
             return self.visit(child)
 
-        def _wrap(child):
+        def _wrap(child, force_use_bracket=False):
             latex = _unwrap(child)
             if isinstance(child, ast.BinOp):
                 cp = priority[type(child.op)] if type(child.op) in priority else 100
                 pp = priority[type(node.op)] if type(node.op) in priority else 100
-                if cp < pp:
+                if cp < pp or (cp == pp and force_use_bracket):
                     return "(" + latex + ")"
             return latex
 
         lhs = node.left
         rhs = node.right
         reprs = {
             ast.Add: (lambda: _wrap(lhs) + " + " + _wrap(rhs)),
-            ast.Sub: (lambda: _wrap(lhs) + " - " + _wrap(rhs)),
+            ast.Sub: (lambda: _wrap(lhs) + " - " + _wrap(rhs, True)),
             ast.Mult: (lambda: _wrap(lhs) + _wrap(rhs)),
             ast.MatMult: (lambda: _wrap(lhs) + _wrap(rhs)),
             ast.Div: (lambda: r"\frac{" + _unwrap(lhs) + "}{" + _unwrap(rhs) + "}"),