How to change the traceback of a custom Exception in Python?

Question

I'm trying to create a custom exception that is raised inside of a class. However, whenever I raise it, the traceback leads to the original 'raise' statement.

Here's a simplified example of what I'm trying to do

# 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

And this is the output

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

Instead of the traceback leading to line 8, I want it to only lead to line 14 like this:

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

Is this possible?

I tried googling it and didn't find anything helpful.

Edit: I would also like to get rid of the __main__. in the last line of the stack trace if possible.

Answer 1

I found an answer (though not very elegant)

try:
    raise ImplementationError("Expected interface")
except ImplementationError as e:
    frame = e.__traceback__.tb_frame.f_back
    tb = e.__traceback__.__class__(None, frame, frame.f_lasti, frame.f_lineno)

    traceback.print_exception(e, e, tb)
    exit()