cleanup, add test

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/pipelines/sequential/ast_helpers.py

@@ -88,7 +88,6 @@ def append_autowrap_source_on_fail():
         _exc_type, _exc_value, exc_tb = sys.exc_info()
         tb_list = traceback.extract_tb(exc_tb)
 
-        collected_sources = []
         for frame in reversed(tb_list):
             if "Autowrapped" in frame.filename:
                 source_lines = linecache.getlines(frame.filename)
@@ -100,12 +99,9 @@ def append_autowrap_source_on_fail():
                     for i, line in enumerate(source_lines)
                 ]
 
-                collected_sources.append(
-                    f"\n--- Autowrapped source for {frame.filename}:{lineno} ---\n"
-                    + "".join(source_lines)
-                )
-
-                new_message = f"{exception}\n\n" + "\n".join(collected_sources)
-                raise RuntimeError(new_message) from exception
+                message = f"{exception}\n\n"
+                message += f"\n--- {frame.filename}:{lineno} ---\n"
+                message += "".join(source_lines)
+                raise RuntimeError(message) from exception
 
         raise exception

tests/llmcompressor/pipelines/sequential/ast_utils.py/test_auto_wrapper.py

@@ -21,13 +21,14 @@ def check_wrapping(
 
     wrapped_lines = ast.unparse(wrapped).splitlines()
     output_lines = textwrap.dedent(output).splitlines()[1:]
+    lines = ("\n".join(wrapped_lines), "\n".join(output_lines))
 
-    assert len(wrapped_lines) == len(output_lines)
+    assert len(wrapped_lines) == len(output_lines), lines
     for wrapped_line, output_line in zip(wrapped_lines, output_lines):
         if "# skip" in output:
             continue
 
-        assert wrapped_line == output_line
+        assert wrapped_line == output_line, lines
 
 
 def test_static_if():
@@ -189,3 +190,24 @@ def forward(a, *b, c=5, **d):
         () = wrapped_0(a, b, c, d)
     """
     check_wrapping(source, output)
+
+
+def test_walrus():
+    """Checks for handling variadic names created via function def"""
+
+    source = """
+    def forward():
+        if (asdf := (1 + 2)):
+            pass
+    """
+    output = """
+    @torch.fx.wrap
+    def wrapped_0():
+        if (asdf := (1 + 2)):
+            pass
+        return (asdf,)
+    
+    def forward():
+        (asdf,) = wrapped_0()
+    """
+    check_wrapping(source, output)