如何精准判断Python类的__add__/__radd__方法是否被sum内置函数调用？

当自定义Python对象实现了__add__和__radd__方法时，我们需要精准判断这些方法是否是在sum内置函数的执行上下文中被调用。现有通过检测代码上下文字符串的方式存在缺陷——比如在sum([a,b]) +5的场景中，最后一次__add__是在sum外部执行，但原逻辑会误判为True。

原朴素示例代码 #

import inspect
import shlex

class Test:
    def __init__(self, val):
        self.val = val
        
    def __add__(self, other):
        curframe = inspect.currentframe()   
        context = inspect.getouterframes(curframe)[1].code_context[0]
        is_summed = 'sum' in list(shlex.shlex(context))
        print('__add__ called from sum():  ', is_summed)       

        try:
            return Test(self.val + other.val)
        except:
            return Test(self.val + other)
    
    def __radd__(self, other):
        curframe = inspect.currentframe()   
        context = inspect.getouterframes(curframe)[1].code_context[0]
        is_summed = 'sum' in list(shlex.shlex(context))
        print('__radd__ called from sum(): ', is_summed)
        
        if other == 0:
            return self
        else:
            return self.__add__(other)
          

if __name__ == '__main__':
    
    a = Test(5)
    b = Test(10)

    print('example 1 - plain add: a + b')
    a + b

    print('\nexample 2 - using sum: sum([a, b])')
    sum([a, b])
    
    print('\nexample 3 - (broken) using sum and an add:sum([a, b]) + 5')
    sum([a, b]) + 5

原执行结果 #

example 1 - plain add: a + b
__add__  called from sum():  False

example 2 - using sum: sum([a, b])
__radd__ called from sum():  True
__add__  called from sum():  True

example 3 - (broken) using sum and an add:sum([a, b]) + 5
__radd__ called from sum():  True
__add__  called from sum():  True
__add__  called from sum():  True # desired result here is False because called outside of "sum"

解决方案：通过调用栈检测sum的执行上下文 #

我们可以利用inspect模块遍历调用栈，检查是否有栈帧来自内置sum函数的执行过程。核心逻辑是：遍历调用栈中的每个帧，判断帧对应的函数名是否为'sum'，且该函数属于builtins模块（避免误判同名自定义函数）。

修改后的代码 #

import inspect
import builtins

class Test:
    def __init__(self, val):
        self.val = val
        
    def _is_called_from_sum(self):
        # 遍历调用栈，检查是否有帧来自内置sum函数
        curframe = inspect.currentframe()
        while curframe:
            frame_func = curframe.f_code.co_name
            # 检查函数名是否为sum，且该函数是内置的sum
            if frame_func == 'sum' and inspect.getmodule(curframe) is builtins:
                return True
            curframe = curframe.f_back
        return False
        
    def __add__(self, other):
        is_summed = self._is_called_from_sum()
        print('__add__ called from sum():  ', is_summed)       

        try:
            return Test(self.val + other.val)
        except:
            return Test(self.val + other)
    
    def __radd__(self, other):
        is_summed = self._is_called_from_sum()
        print('__radd__ called from sum(): ', is_summed)
        
        if other == 0:
            return self
        else:
            return self.__add__(other)
          

if __name__ == '__main__':
    
    a = Test(5)
    b = Test(10)

    print('example 1 - plain add: a + b')
    a + b

    print('\nexample 2 - using sum: sum([a, b])')
    sum([a, b])
    
    print('\nexample 3 - fixed: using sum and an add:sum([a, b]) + 5')
    sum([a, b]) + 5

修改后的执行结果 #

example 1 - plain add: a + b
__add__  called from sum():  False

example 2 - using sum: sum([a, b])
__radd__ called from sum():  True
__add__  called from sum():  True

example 3 - fixed: using sum and an add:sum([a, b]) + 5
__radd__ called from sum():  True
__add__  called from sum():  True
__add__  called from sum():  False  # 符合预期，此时__add__在sum外部调用

说明 #

• _is_called_from_sum方法负责遍历调用栈：从当前帧开始，逐层向上检查每个栈帧的函数名和所属模块，确认是否来自内置sum函数。
• 这种方式能精准区分sum内部和外部的调用场景，解决了原方案中通过字符串匹配导致的误判问题。

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