Make srepr work correctly on Add and Mul subclasses

asmeurer · asmeurer · commit 49e52daf1342 · 2019-03-29T15:36:32.000-06:00
Fixes sympy#16430.
diff --git a/sympy/printing/repr.py b/sympy/printing/repr.py
@@ -48,9 +48,10 @@ def _print_Add(self, expr, order=None):
         args = self._as_ordered_terms(expr, order=order)
         nargs = len(args)
         args = map(self._print, args)
+        clsname = type(expr).__name__
         if nargs > 255:  # Issue #10259, Python < 3.7
-            return "Add(*[%s])" % ", ".join(args)
-        return "Add(%s)" % ", ".join(args)
+            return clsname + "(*[%s])" % ", ".join(args)
+        return clsname + "(%s)" % ", ".join(args)
 
     def _print_Cycle(self, expr):
         return expr.__repr__()
@@ -138,9 +139,10 @@ def _print_Mul(self, expr, order=None):
 
         nargs = len(args)
         args = map(self._print, args)
+        clsname = type(expr).__name__
         if nargs > 255:  # Issue #10259, Python < 3.7
-            return "Mul(*[%s])" % ", ".join(args)
-        return "Mul(%s)" % ", ".join(args)
+            return clsname + "(*[%s])" % ", ".join(args)
+        return clsname + "(%s)" % ", ".join(args)
 
     def _print_Rational(self, expr):
         return 'Rational(%s, %s)' % (self._print(expr.p), self._print(expr.q))
diff --git a/sympy/printing/tests/test_repr.py b/sympy/printing/tests/test_repr.py
@@ -1,7 +1,7 @@
 from sympy.utilities.pytest import raises
 from sympy import (symbols, Function, Integer, Matrix, Abs,
     Rational, Float, S, WildFunction, ImmutableDenseMatrix, sin, true, false, ones,
-    sqrt, root, AlgebraicNumber, Symbol, Dummy, Wild)
+    sqrt, root, AlgebraicNumber, Symbol, Dummy, Wild, MatrixSymbol)
 from sympy.core.compatibility import exec_
 from sympy.geometry import Point, Ellipse
 from sympy.printing import srepr
@@ -274,3 +274,11 @@ def test_Naturals0():
 
 def test_Reals():
     sT(S.Reals, "Reals")
+
+def test_matrix_expressions():
+    n = symbols('n', integer=True)
+    A = MatrixSymbol("A", n, n)
+    B = MatrixSymbol("B", n, n)
+    sT(A, "MatrixSymbol(Symbol('A'), Symbol('n', integer=True), Symbol('n', integer=True))")
+    sT(A*B, "MatMul(MatrixSymbol(Symbol('A'), Symbol('n', integer=True), Symbol('n', integer=True)), MatrixSymbol(Symbol('B'), Symbol('n', integer=True), Symbol('n', integer=True)))")
+    sT(A + B, "MatAdd(MatrixSymbol(Symbol('A'), Symbol('n', integer=True), Symbol('n', integer=True)), MatrixSymbol(Symbol('B'), Symbol('n', integer=True), Symbol('n', integer=True)))")
