如何修改Python自定义Exception的回溯信息并移除__main__前缀？

尝试在类中定义并抛出自定义Exception，但抛出时的回溯信息会指向装饰器wrapper函数内的原始raise语句。简化示例代码如下：

# Errors
class NotImplementedError(Exception):
    def __init__(self, msg=""):
        super().__init__(msg)

class ImplementationError(Exception):
    def __init__(self, msg):
        super().__init__(msg)
        
# Interface decorator class
class Interface:
    def __init__(self, cls):
        self.cls = cls

    def __call__(self):
        return self.cls()

def implement(interface):
    if type(interface) != Interface:
        raise ImplementationError("Expected interface")
    
    interface = interface.cls

    def wrapper(cls):
        iKeys = interface.__dict__.keys()
        cKeys = cls.__dict__.keys()

        for iKey in iKeys:
            if iKey not in cKeys:
                msg = (f"'{iKey}' not implemented in {cls.__name__}")
                raise NotImplementedError(msg)
        
        return cls

    return wrapper

@Interface
class Shape:
    origin = None

    def area(self): pass
    def bounding_rect(self): pass

@implement(Shape)
class Square:
    pass

当前输出的回溯信息：

Traceback (most recent call last):
  File "main.py", line 48, in <module>
    class Square:
  File "main.py", line 34, in wrapper
    raise NotImplementedError(msg)
__main__.NotImplementedError: 'origin' not implemented in Square

期望的回溯信息：

Traceback (most recent call last):
  File "main.py", line 48, in <module>
    class Square:
NotImplementedError: 'origin' not implemented in Square

需求：

• 回溯信息仅指向定义Square类的行
• 移除栈追踪最后一行的__main__前缀

1. 调整回溯信息，移除wrapper栈帧 #

Python允许手动修改异常的__traceback__属性，跳过装饰器wrapper对应的栈帧。修改implement装饰器中的wrapper函数，先导入sys模块，再在抛出异常时调整traceback：

import sys

# ... 其他代码不变 ...

def implement(interface):
    if type(interface) != Interface:
        raise ImplementationError("Expected interface")
    
    interface = interface.cls

    def wrapper(cls):
        iKeys = interface.__dict__.keys()
        cKeys = cls.__dict__.keys()

        for iKey in iKeys:
            if iKey not in cKeys:
                msg = (f"'{iKey}' not implemented in {cls.__name__}")
                exc = NotImplementedError(msg)
                # 跳过wrapper函数的栈帧
                exc.__traceback__ = exc.__traceback__.tb_next
                raise exc
        
        return cls

    return wrapper

exc.__traceback__.tb_next会跳过当前wrapper函数对应的栈帧，让回溯直接指向类定义的行。

2. 移除异常的__main__前缀 #

通过修改自定义异常类的__module__属性为空字符串，Python在打印异常时就不会显示模块前缀：

# Errors
class NotImplementedError(Exception):
    def __init__(self, msg=""):
        super().__init__(msg)
        self.__module__ = ""  # 移除模块前缀

class ImplementationError(Exception):
    def __init__(self, msg):
        super().__init__(msg)
        self.__module__ = ""  # 移除模块前缀

修改后的完整代码 #

import sys

# Errors
class NotImplementedError(Exception):
    def __init__(self, msg=""):
        super().__init__(msg)
        self.__module__ = ""

class ImplementationError(Exception):
    def __init__(self, msg):
        super().__init__(msg)
        self.__module__ = ""
        
# Interface decorator class
class Interface:
    def __init__(self, cls):
        self.cls = cls

    def __call__(self):
        return self.cls()

def implement(interface):
    if type(interface) != Interface:
        raise ImplementationError("Expected interface")
    
    interface = interface.cls

    def wrapper(cls):
        iKeys = interface.__dict__.keys()
        cKeys = cls.__dict__.keys()

        for iKey in iKeys:
            if iKey not in cKeys:
                msg = (f"'{iKey}' not implemented in {cls.__name__}")
                exc = NotImplementedError(msg)
                exc.__traceback__ = exc.__traceback__.tb_next
                raise exc
        
        return cls

    return wrapper

@Interface
class Shape:
    origin = None

    def area(self): pass
    def bounding_rect(self): pass

@implement(Shape)
class Square:
    pass

运行后输出的回溯信息将完全符合期望：

Traceback (most recent call last):
  File "main.py", line 51, in <module>
    class Square:
NotImplementedError: 'origin' not implemented in Square

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