diff --git a/numba_cuda/numba/cuda/misc/firstlinefinder.py b/numba_cuda/numba/cuda/misc/firstlinefinder.py index de77112c6..16daf889a 100644 --- a/numba_cuda/numba/cuda/misc/firstlinefinder.py +++ b/numba_cuda/numba/cuda/misc/firstlinefinder.py @@ -7,6 +7,8 @@ """ import ast +import inspect +import textwrap class FindDefFirstLine(ast.NodeVisitor): @@ -16,17 +18,23 @@ class FindDefFirstLine(ast.NodeVisitor): first_stmt_line : int or None This stores the first statement line number if the definition is found. Or, ``None`` if the definition is not found. + def_lineno : int or None + This stores the 'def' line number if the definition is found. + Or, ``None`` if the definition is not found. """ - def __init__(self, code): + def __init__(self, name, firstlineno): """ Parameters ---------- - code : - The function's code object. + name : str + The function's name (co_name). + firstlineno : int + The function's first line number (co_firstlineno), adjusted for + any offset if the source is a fragment. """ - self._co_name = code.co_name - self._co_firstlineno = code.co_firstlineno + self._co_name = name + self._co_firstlineno = firstlineno self.first_stmt_line = None self.def_lineno = None @@ -76,8 +84,11 @@ def _is_docstring(node): def get_func_body_first_lineno(pyfunc): """ - Look up the first line of function body using the file in - ``pyfunc.__code__.co_filename``. + Look up the first line of function body. + + Uses inspect.getsourcelines() which works for both regular .py files + (via linecache reading from disk) and Jupyter notebook cells (via + IPython's linecache registration). Returns ------- @@ -87,21 +98,29 @@ def get_func_body_first_lineno(pyfunc): """ co = pyfunc.__code__ try: - with open(co.co_filename) as fin: - file_content = fin.read() - except (FileNotFoundError, OSError): - return + lines, offset = inspect.getsourcelines(pyfunc) + source = "".join(lines) + offset = offset - 1 + except (OSError, TypeError): + return None + + tree = ast.parse(textwrap.dedent(source)) + finder = FindDefFirstLine(co.co_name, co.co_firstlineno - offset) + finder.visit(tree) + if finder.first_stmt_line: + return finder.first_stmt_line + offset else: - tree = ast.parse(file_content) - finder = FindDefFirstLine(co) - finder.visit(tree) - return finder.first_stmt_line + # No first line found. + return None def get_func_def_lineno(pyfunc): """ - Look up the line number of the function definition ('def' line) using - the file in ``pyfunc.__code__.co_filename``. + Look up the line number of the function definition ('def' line). + + Uses inspect.getsourcelines() which works for both regular .py files + (via linecache reading from disk) and Jupyter notebook cells (via + IPython's linecache registration). Returns ------- @@ -111,12 +130,17 @@ def get_func_def_lineno(pyfunc): """ co = pyfunc.__code__ try: - with open(co.co_filename) as fin: - file_content = fin.read() - except (FileNotFoundError, OSError): - return + lines, offset = inspect.getsourcelines(pyfunc) + source = "".join(lines) + offset = offset - 1 + except (OSError, TypeError): + return None + + tree = ast.parse(textwrap.dedent(source)) + finder = FindDefFirstLine(co.co_name, co.co_firstlineno - offset) + finder.visit(tree) + if finder.def_lineno: + return finder.def_lineno + offset else: - tree = ast.parse(file_content) - finder = FindDefFirstLine(co) - finder.visit(tree) - return finder.def_lineno + # No def line found. + return None diff --git a/numba_cuda/numba/cuda/tests/cudapy/test_debuginfo.py b/numba_cuda/numba/cuda/tests/cudapy/test_debuginfo.py index 4f1388424..721eb8003 100644 --- a/numba_cuda/numba/cuda/tests/cudapy/test_debuginfo.py +++ b/numba_cuda/numba/cuda/tests/cudapy/test_debuginfo.py @@ -626,6 +626,65 @@ def foo(): # and refers to the offending function self.assertIn(str(foo.py_func), msg) + def test_linecache_source(self): + """Test that source from linecache (like Jupyter notebooks) works. + + This simulates how Jupyter/IPython registers cell source in linecache, + allowing inspect.getsourcelines() to find it even though the file + doesn't exist on disk. Fixes issue #721. + """ + import linecache + + # Source with a multi-line decorator + strsrc = dedent(""" + @cuda.jit( + "void(int32[:])", + debug=True, + opt=False + ) + def foo(x): + x[0] = 1 + """).strip() + + # Simulate Jupyter by registering source in linecache + fake_filename = "" + lines = [line + "\n" for line in strsrc.splitlines()] + linecache.cache[fake_filename] = ( + len(strsrc), + None, # mtime=None means never expire + lines, + fake_filename, + ) + + try: + # Compile and execute using the fake filename + code = compile(strsrc, fake_filename, "exec") + exec_globals = {"cuda": cuda} + exec(code, exec_globals) + foo = exec_globals["foo"] + + # Should NOT produce a warning since source is in linecache + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always", NumbaDebugInfoWarning) + ignore_internal_warnings() + foo[1, 1](cuda.to_device(np.zeros(1, dtype=np.int32))) + + # Filter for NumbaDebugInfoWarning specifically + debug_warnings = [ + x for x in w if x.category == NumbaDebugInfoWarning + ] + self.assertEqual( + len(debug_warnings), + 0, + msg=f"Unexpected warning: {debug_warnings}", + ) + + # Verify debug info is present in the PTX + self._check(foo, sig=(types.int32[:],), expect=True) + finally: + # Clean up linecache + linecache.cache.pop(fake_filename, None) + def test_no_if_op_bools_declared(self): @cuda.jit( "int64(boolean, boolean)",