17759: convenience class symbolic ExpressionTreeWalker(Converter)

Author: rwst

src/sage/symbolic/expression_conversions.py

@@ -1672,80 +1672,136 @@ def composition(self, ex, operator):
         else:
             return res
 
-class SubstituteFunction(Converter):
-    def __init__(self, ex, original, new):
+
+class ExpressionTreeWalker(Converter):
+    def __init__(self, ex):
         """
-        A class that walks the tree and replaces occurrences of a
-        function with another.
+        A class that walks the tree. Mainly for subclassing.
 
         EXAMPLES::
 
-            sage: from sage.symbolic.expression_conversions import SubstituteFunction
-            sage: foo = function('foo'); bar = function('bar')
-            sage: s = SubstituteFunction(foo(x), foo, bar)
-            sage: s(1/foo(foo(x)) + foo(2))
-            1/bar(bar(x)) + bar(2)
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: from sage.symbolic.random_tests import random_expr
+            sage: ex = sin(atan(0,hold=True)+hypergeometric((1,),(1,),x))
+            sage: s = ExpressionTreeWalker(ex)
+            sage: bool(s() == ex)
+            True
+            sage: foo = random_expr(20, nvars=2)
+            sage: s = ExpressionTreeWalker(foo)
+            sage: bool(s() == foo)
+            True
         """
-        self.original = original
-        self.new = new
         self.ex = ex
 
     def symbol(self, ex):
         """
         EXAMPLES::
 
-            sage: from sage.symbolic.expression_conversions import SubstituteFunction
-            sage: foo = function('foo'); bar = function('bar')
-            sage: s = SubstituteFunction(foo(x), foo, bar)
-            sage: s.symbol(x)
-            x
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: s = ExpressionTreeWalker(x)
+            sage: bool(s.symbol(x) == x)
+            True
         """
         return ex
 
     def pyobject(self, ex, obj):
         """
         EXAMPLES::
 
-            sage: from sage.symbolic.expression_conversions import SubstituteFunction
-            sage: foo = function('foo'); bar = function('bar')
-            sage: s = SubstituteFunction(foo(x), foo, bar)
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
             sage: f = SR(2)
-            sage: s.pyobject(f, f.pyobject())
-            2
-            sage: _.parent()
-            Symbolic Ring
+            sage: s = ExpressionTreeWalker(f)
+            sage: bool(s.pyobject(f, f.pyobject()) == f.pyobject())
+            True
         """
         return ex
 
     def relation(self, ex, operator):
         """
         EXAMPLES::
 
-            sage: from sage.symbolic.expression_conversions import SubstituteFunction
-            sage: foo = function('foo'); bar = function('bar')
-            sage: s = SubstituteFunction(foo(x), foo, bar)
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: foo = function('foo')
             sage: eq = foo(x) == x
-            sage: s.relation(eq, eq.operator())
-            bar(x) == x
+            sage: s = ExpressionTreeWalker(eq)
+            sage: s.relation(eq, eq.operator()) == eq
+            True
         """
         return operator(self(ex.lhs()), self(ex.rhs()))
 
     def arithmetic(self, ex, operator):
         """
+        EXAMPLES::
+
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: foo = function('foo')
+            sage: f = x*foo(x) + pi/foo(x)
+            sage: s = ExpressionTreeWalker(f)
+            sage: bool(s.arithmetic(f, f.operator()) == f)
+            True
+        """
+        return reduce(operator, map(self, ex.operands()))
+
+    def composition(self, ex, operator):
+        """
+        EXAMPLES::
+
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: foo = function('foo')
+            sage: f = foo(atan2(0, 0, hold=True))
+            sage: s = ExpressionTreeWalker(f)
+            sage: bool(s.composition(f, f.operator()) == f)
+            True
+        """
+        from sage.symbolic.function import Function
+        if isinstance(operator, Function):
+            return operator(*map(self, ex.operands()), hold=True)
+        else:
+            return operator(*map(self, ex.operands()))
+
+    def derivative(self, ex, operator):
+        """
+        EXAMPLES::
+
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: foo = function('foo')
+            sage: f = foo(x).diff(x)
+            sage: s = ExpressionTreeWalker(f)
+            sage: bool(s.derivative(f, f.operator()) == f)
+            True
+        """
+        return operator(*map(self, ex.operands()))
+
+    def tuple(self, ex):
+        """
+        EXAMPLES::
+
+            sage: from sage.symbolic.expression_conversions import ExpressionTreeWalker
+            sage: foo = function('foo')
+            sage: f = hypergeometric((1,2,3,),(x,),x)
+            sage: s = ExpressionTreeWalker(f)
+            sage: bool(s() == f)
+            True
+        """
+        return ex.operands()
+
+class SubstituteFunction(ExpressionTreeWalker):
+    def __init__(self, ex, original, new):
+        """
+        A class that walks the tree and replaces occurrences of a
+        function with another.
+
         EXAMPLES::
 
             sage: from sage.symbolic.expression_conversions import SubstituteFunction
             sage: foo = function('foo'); bar = function('bar')
             sage: s = SubstituteFunction(foo(x), foo, bar)
-            sage: f = x*foo(x) + pi/foo(x)
-            sage: s.arithmetic(f, f.operator())
-            x*bar(x) + pi/bar(x)
+            sage: s(1/foo(foo(x)) + foo(2))
+            1/bar(bar(x)) + bar(2)
         """
-        if operator == add_vararg:
-            operator = _operator.add
-        elif operator == mul_vararg:
-            operator = _operator.mul
-        return reduce(operator, map(self, ex.operands()))
+        self.original = original
+        self.new = new
+        self.ex = ex
 
     def composition(self, ex, operator):
         """
@@ -1798,3 +1854,4 @@ def derivative(self, ex, operator):
             return operator.change_function(self.new)(*map(self,ex.operands()))
         else:
             return operator(*map(self, ex.operands()))
+