Improve logic for finding print statements (#58)

* fix logic for finding print statement

* fix test

Author: samuelcolvin

pyproject.toml

@@ -9,7 +9,7 @@ include = ["/README.md", "/Makefile", "/pytest_examples", "/tests"]
 
 [project]
 name = "pytest-examples"
-version = "0.0.16"
+version = "0.0.17"
 description = "Pytest plugin for testing examples in docstrings and markdown files."
 authors = [
     {name = "Samuel Colvin", email = "[email protected]"},

pytest_examples/run_code.py

@@ -156,8 +156,8 @@ def __call__(self, *args: Any, sep: str = ' ', **kwargs: Any) -> None:
         frame = inspect.stack()[parent_frame_id]
 
         if self._include_file(frame, args):
-            # -1 to account for the line number being 1-indexed
-            s = PrintStatement(frame.lineno, sep, [Arg(arg) for arg in args])
+            lineno = self._find_line_number(frame)
+            s = PrintStatement(lineno, sep, [Arg(arg) for arg in args])
             self.statements.append(s)
 
     def _include_file(self, frame: inspect.FrameInfo, args: Sequence[Any]) -> bool:
@@ -166,6 +166,17 @@ def _include_file(self, frame: inspect.FrameInfo, args: Sequence[Any]) -> bool:
         else:
             return self.file.samefile(frame.filename)
 
+    def _find_line_number(self, inspect_frame: inspect.FrameInfo) -> int:
+        """Find the line number of the print statement in the file that is being executed."""
+        frame = inspect_frame.frame
+        while True:
+            if self.file.samefile(frame.f_code.co_filename):
+                return frame.f_lineno
+            elif frame.f_back:
+                frame = frame.f_back
+            else:
+                raise RuntimeError(f'Could not find line number of print statement at {inspect_frame}')
+
 
 class InsertPrintStatements:
     def __init__(
@@ -256,18 +267,6 @@ def _insert_print_args(
 triple_quotes_prefix_re = re.compile('^ *(?:"{3}|\'{3})', re.MULTILINE)
 
 
-def find_print_line(lines: list[str], line_no: int) -> int:
-    """For 3.7 we have to reverse through lines to find the print statement lint."""
-    return line_no
-
-    for back in range(100):
-        new_line_no = line_no - back
-        m = re.search(r'^ *print\(', lines[new_line_no - 1])
-        if m:
-            return new_line_no
-    return line_no
-
-
 def remove_old_print(lines: list[str], line_index: int) -> None:
     """Remove the old print statement."""
     try:
@@ -294,12 +293,12 @@ def remove_old_print(lines: list[str], line_index: int) -> None:
 def find_print_location(example: CodeExample, line_no: int) -> tuple[int, int]:
     """Find the line and column of the print statement.
 
-    :param example: the `CodeExample`
-    :param line_no: The line number on which the print statement starts (or approx on 3.7)
-    :return: tuple if `(line, column)` of the print statement
-    """
-    # For 3.7 we have to reverse through lines to find the print statement lint
+    Args:
+        example: the `CodeExample`
+        line_no: The line number on which the print statement starts or approx
 
+    Return: tuple if `(line, column)` of the print statement
+    """
     m = ast.parse(example.source, filename=example.path.name)
     return find_print(m, line_no) or (line_no, 0)
 

tests/test_insert_print.py

@@ -1,5 +1,7 @@
 from __future__ import annotations as _annotations
 
+import sys
+
 import pytest
 from _pytest.outcomes import Failed
 
@@ -397,8 +399,6 @@ def main():
 
 
 def test_run_main_print(tmp_path, eval_example):
-    # note this file is no written here as it's not required
-    md_file = tmp_path / 'test.md'
     python_code = """
 main_called = False
 
@@ -408,16 +408,15 @@ def main():
     print(1, 2, 3)
     #> 1 2 3
 """
-    example = CodeExample.create(python_code, path=md_file)
+    # note this file is no written here as it's not required
+    example = CodeExample.create(python_code, path=tmp_path / 'test.md')
     eval_example.set_config(line_length=30)
 
     module_dict = eval_example.run_print_check(example, call='main')
     assert module_dict['main_called']
 
 
 def test_run_main_print_async(tmp_path, eval_example):
-    # note this file is no written here as it's not required
-    md_file = tmp_path / 'test.md'
     python_code = """
 main_called = False
 
@@ -427,22 +426,22 @@ async def main():
     print(1, 2, 3)
     #> 1 2 3
 """
-    example = CodeExample.create(python_code, path=md_file)
+    # note this file is no written here as it's not required
+    example = CodeExample.create(python_code, path=tmp_path / 'test.md')
     eval_example.set_config(line_length=30)
 
     module_dict = eval_example.run_print_check(example, call='main')
     assert module_dict['main_called']
 
 
 def test_custom_include_print(tmp_path, eval_example):
-    # note this file is no written here as it's not required
-    md_file = tmp_path / 'test.md'
     python_code = """
 print('yes')
 #> yes
 print('no')
 """
-    example = CodeExample.create(python_code, path=md_file)
+    # note this file is no written here as it's not required
+    example = CodeExample.create(python_code, path=tmp_path / 'test.md')
     eval_example.set_config(line_length=30)
 
     def custom_include_print(path, frame, args):
@@ -451,3 +450,29 @@ def custom_include_print(path, frame, args):
     eval_example.include_print = custom_include_print
 
     eval_example.run_print_check(example, call='main')
+
+
+def test_print_different_file(tmp_path, eval_example):
+    other_file = tmp_path / 'other.py'
+    other_code = """
+def does_print():
+    print('hello')
+    """
+    other_file.write_text(other_code)
+    sys.path.append(str(tmp_path))
+    python_code = """
+import other
+
+other.does_print()
+#> hello
+"""
+    example = CodeExample.create(python_code, path=tmp_path / 'test.md')
+
+    eval_example.include_print = lambda p, f, a: True
+
+    eval_example.run_print_check(example, call='main')
+
+    del sys.modules['other']
+    other_file.write_text(('\n' * 30) + other_code)
+
+    eval_example.run_print_check(example, call='main')