How to capture the traceback for an error in python along with source code?

When I run a vscode notebook cell or a python file with these contents:

def foo():
    return 4/0  # This will cause a division by zero exception

foo()

I get the following error traceback which is very meaningful and tells me at what line along with the line of code that errors:

---------------------------------------------------------------------------
ZeroDivisionError                         Traceback (most recent call last)
Cell In[12], line 4
      1 def foo():
      2     return 4/0  # This will cause a division by zero exception
----> 4 foo()

Cell In[12], line 2, in foo()
      1 def foo():
----> 2     return 4/0

ZeroDivisionError: division by zero

However when I run the program using a exec block I am unable to see the source code. How do I capture this?

import traceback
import sys

code = """def foo():
    return 4/0  # This will cause a division by zero exception

foo()
"""

try:
    exec(code)
except Exception as e:
    traceback.print_exc() 

gives:

Traceback (most recent call last):
  File "/var/folders/fd/f7mj1ws56pq6xvyzj4zv1r_m0000gn/T/ipykernel_39670/712868849.py", line 11, in <module>
    exec(code)
  File "<string>", line 4, in <module>
  File "<string>", line 2, in foo
ZeroDivisionError: division by zero

The traceback printing mechanism relies on being able to read the line from disk (using the linecache module).

For exec'd code, there is no disk (after all, that's why it's saying File "<string>", line 2).

Looking at what traceback is doing you can patch linecache.getline to do a special thing for <string>s:

import traceback
import linecache
import unittest.mock

code = """def foo():
    return 4/0  # This will cause a division by zero exception

foo()
"""

orig_getline = linecache.getline


def new_getline(filename, lineno, *args, **kwargs):
    if filename == "<string>":
        return code.splitlines()[lineno - 1]
    return orig_getline(filename, lineno, *args, **kwargs)


with unittest.mock.patch("linecache.getline", new_getline):
    try:
        exec(code)
    except Exception as e:
        traceback.print_exc()

et voilà:

Traceback (most recent call last):
  File "F:\build\so-misc\so76878031.py", line 23, in <module>
    exec(code)
  File "<string>", line 4, in <module>
    foo()
  File "<string>", line 2, in foo
    return 4/0  # This will cause a division by zero exception
ZeroDivisionError: division by zero