From 29c69c44165320df7d757e0e51f4a21c54bb592f Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Mon, 24 Mar 2025 11:42:19 +0000
Subject: [PATCH] [mypyc] Display IR on annotate test failure

This makes it easier to figure out why a test is failing.

Example output on failure:
```
=================================== FAILURES ====================================
________________________ testAnnotateTwoOperationsOnLine ________________________
data: /Users/jukka/src/mypy/mypyc/test-data/annotate-basic.test:18:
Failed: Invalid source code output (/Users/jukka/src/mypy/mypyc/test-data/annotate-basic.test, line 18)
----------------------------- Captured stdout call ------------------------------
Generated IR:

def f(x):
    x :: object
    r0 :: str
    r1, r2, r3 :: object
L0:
    r0 = 'foo'
    r1 = CPyObject_GetAttr(x, r0)
    r2 = object 1
    r3 = PyNumber_Add(r1, r2)
    return r3
----------------------------- Captured stderr call ------------------------------
Expected:
  main:2: Get non-native attribute "foo". Generic "+" operation.x (diff)
Actual:
  main:2: Get non-native attribute "foo". Generic "+" operation. (diff)

Alignment of first line difference:
  E: ...Generic "+" operation.x
  A: ...Generic "+" operation.
                              ^
Update the test output using --update-data (implies -n0; you can additionally use the -k selector to update only specific tests)
```
---
 mypyc/test/test_annotate.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/mypyc/test/test_annotate.py b/mypyc/test/test_annotate.py
index 40b28195b5a5..f429fb28cd55 100644
--- a/mypyc/test/test_annotate.py
+++ b/mypyc/test/test_annotate.py
@@ -8,6 +8,7 @@
 from mypy.test.config import test_temp_dir
 from mypy.test.data import DataDrivenTestCase
 from mypyc.annotate import generate_annotations
+from mypyc.ir.pprint import format_func
 from mypyc.test.testutil import (
     ICODE_GEN_BUILTINS,
     MypycDataSuite,
@@ -39,8 +40,9 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
             for i, line in enumerate(testcase.input):
                 if "# A:" in line:
                     msg = line.rpartition("# A:")[2].strip()
-                    expected_output.append(f"{i + 1}: {msg}")
+                    expected_output.append(f"main:{i + 1}: {msg}")
 
+            ir = None
             try:
                 ir, tree = build_ir_for_single_file2(testcase.input, options)
             except CompileError as e:
@@ -50,6 +52,16 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
                 actual = []
                 for line_num, line_anns in annotations.annotations.items():
                     s = " ".join(line_anns)
-                    actual.append(f"{line_num}: {s}")
+                    actual.append(f"main:{line_num}: {s}")
 
-            assert_test_output(testcase, actual, "Invalid source code output", expected_output)
+            try:
+                assert_test_output(testcase, actual, "Invalid source code output", expected_output)
+            except BaseException:
+                if ir:
+                    print("Generated IR:\n")
+                    for fn in ir.functions:
+                        if fn.name == "__top_level__":
+                            continue
+                        for s in format_func(fn):
+                            print(s)
+                raise