自定义SymPy Function中__str__方法的行为问题及解决需求

问题背景 #

自定义SymPy Function子类时，希望其在转换为字符串（包括作为表达式一部分）时，能像_latex方法一样实现自定义输出。现有代码中，单独打印实例时__str__生效，但实例作为表达式一部分时，字符串输出仍为默认格式。

现有代码 #

from sympy import Function, Symbol, latex
class TestClass(Function):
    def _latex(self, printer):
        return 'This is a latex test!'
    def __str__(self):
        return 'This is a str test!'

a = TestClass(Symbol('a'))
ab = TestClass(Symbol('a')) + TestClass(Symbol('b'))

测试结果 #

执行以下代码：

print(latex(ab))
print(str(a))

输出符合预期：

This is a latex test! + This is a latex test!
This is a str test!

但执行print(str(ab))时，输出为默认格式：

TestClass(a) + TestClass(b)

解决方案 #

SymPy的表达式字符串打印由StrPrinter类处理，默认不会调用自定义的__str__方法。要实现自定义行为，可采用以下两种方法：

方法1：在类中定义_str方法 #

StrPrinter会优先调用对象的_str方法（如果存在），只需在TestClass中添加该方法即可：

from sympy import Function, Symbol, latex
class TestClass(Function):
    def _latex(self, printer):
        return 'This is a latex test!'
    def _str(self, printer):
        return 'This is a str test!'
    def __str__(self):
        return self._str(None)

a = TestClass(Symbol('a'))
ab = TestClass(Symbol('a')) + TestClass(Symbol('b'))

# 测试输出
print(latex(ab))
print(str(a))
print(str(ab))

输出结果：

This is a latex test! + This is a latex test!
This is a str test!
This is a str test! + This is a str test!

方法2：为StrPrinter注册自定义打印函数 #

通过为StrPrinter添加针对TestClass的处理函数，实现自定义打印逻辑：

from sympy import Function, Symbol, latex, StrPrinter
class TestClass(Function):
    def _latex(self, printer):
        return 'This is a latex test!'
    def __str__(self):
        return 'This is a str test!'

# 注册TestClass的字符串打印逻辑
def print_testclass(printer, expr):
    return expr.__str__()

StrPrinter._print_TestClass = print_testclass

a = TestClass(Symbol('a'))
ab = TestClass(Symbol('a')) + TestClass(Symbol('b'))

# 测试输出
print(latex(ab))
print(str(a))
print(str(ab))

该方法同样能得到预期输出，适合需要为多个类扩展打印逻辑的场景。

内容的提问来源于stack exchange，提问作者Michael Gfrerer