Fix IndexError in traceback formatter

This protects us from differences between the file on disk and the file
that was originally loaded.

This also corrects a bug where we'd incorrectly show the last line of
the file if the detected line number was 0 (as it is for shim frames).

Signed-off-by: Kirill Ignatev <kiri11@users.noreply.github.com>

Author: kiri11

news/191.bugfix.rst

@@ -0,0 +1 @@
+Fix ``IndexError`` in traceback formatting when the source file has fewer lines than expected.

src/pystack/traceback_formatter.py

@@ -27,18 +27,19 @@ def format_frame(frame: PyFrame) -> Iterable[str]:
     if os.path.exists(code.filename):
         with open(code.filename, "r") as fp:
             lines = fp.readlines()
-        source = lines[code.location.lineno - 1]
-        line_start, line_end, col_start, col_end = code.location
-        if col_start == col_end == 0:
-            yield f"        {source.strip()}"
-        else:
-            if line_end != line_start:
-                col_end = len(source)
-            a = source[:col_start]
-            b = source[col_start:col_end].strip("\n")
-            c = source[col_end:]
-            final = f'{a}{colored(b, color="blue")}{c}'
-            yield f"        {final.strip()}"
+        if 1 <= code.location.lineno <= len(lines):
+            source = lines[code.location.lineno - 1]
+            line_start, line_end, col_start, col_end = code.location
+            if col_start == col_end == 0:
+                yield f"        {source.strip()}"
+            else:
+                if line_end != line_start:
+                    col_end = len(source)
+                a = source[:col_start]
+                b = source[col_start:col_end].strip("\n")
+                c = source[col_end:]
+                final = f'{a}{colored(b, color="blue")}{c}'
+                yield f"        {final.strip()}"
 
     if frame.arguments:
         yield f"      {colored('Arguments:', attrs=['faint'])}"

tests/unit/test_traceback_formatter.py

@@ -241,6 +241,53 @@ def test_traceback_formatter_no_mergeable_native_frames():
     ]
 
 
+def test_traceback_formatter_with_source_lineno_out_of_range():
+    # GIVEN
+
+    frame = PyFrame(
+        prev=None,
+        next=None,
+        code=PyCodeObject(
+            filename="template.html",
+            scope="root",
+            location=LocationInfo(13, 13, 0, 0),
+        ),
+        arguments={"arg": "value"},
+        locals={"local": "value"},
+        is_entry=True,
+        is_shim=False,
+    )
+
+    thread = PyThread(
+        tid=1,
+        frame=frame,
+        native_frames=[],
+        holds_the_gil=False,
+        is_gc_collecting=False,
+        python_version=(3, 8),
+    )
+
+    # WHEN
+    source_data = ""
+    with (
+        patch("builtins.open", mock_open(read_data=source_data)),
+        patch("os.path.exists", return_value=True),
+    ):
+        lines = list(format_thread(thread, NativeReportingMode.OFF))
+
+    # THEN
+
+    assert lines == [
+        "Traceback for thread 1 [] (most recent call last):",
+        '    (Python) File "template.html", line 13, in root',
+        "      Arguments:",
+        "        arg: value",
+        "      Locals:",
+        "        local: value",
+        "",
+    ]
+
+
 def test_traceback_formatter_with_source():
     # GIVEN
 
@@ -1621,12 +1668,10 @@ def test_native_traceback_with_shim_frames():
         '    (Python) File "file1.py", line 1, in function1',
         '        x = "This is the line 1" or (1+1)',
         '    (Python) File "<shim>", line 0, in <shim>',
-        '        x = "This is the line 4" or (1+1)',
         '    (C) File "native_file2.c", line 2, in native_function2 (library.so)',
         '    (Python) File "file2.py", line 2, in function2',
         '        x = "This is the line 2" or (1+1)',
         '    (Python) File "<shim>", line 0, in <shim>',
-        '        x = "This is the line 4" or (1+1)',
         '    (C) File "native_file3.c", line 3, in native_function3 (library.so)',
         '    (Python) File "file3.py", line 3, in function3',
         '        x = "This is the line 3" or (1+1)',