feat(arxjit): add source extraction for decorated functions

[Jaskirat-s7]

Summary

Sprint 2, source extraction (wiki "Proposed Issue 3: Implement Python function source extraction"). This is the second of the four small Sprint 2 PRs agreed on Discord, building on the diagnostics/error layer merged in #95.

It adds arxjit.source.extract_source(), the first stage of the @jit pipeline: given the decorated Python function, produce a decorator-free source string and a parsed ast function node that the upcoming validation pass (and later the Python-AST -> ASTx lowering) can consume.

Updated after review (6e93497): the original version dedented the block with textwrap.dedent before parsing. Per the review findings, the source text is now never modified — indented definitions are parsed inside a synthetic if True: wrapper that preserves lexical content, and __wrapped__ chains are unwrapped before filename attribution. ExtractedSource.source is therefore verbatim file text (not a normalized representation, and for nested functions not standalone-parseable); the parsed node is the canonical artifact. A normalized source representation for the cache-key work is tracked in #97.

Issue 3 checklist coverage:

• Retrieve the source code of decorated Python functions — inspect.getsourcelines (follows __wrapped__, so an actual JitFunction resolves to the function it wraps)
• Normalize indentation — reinterpreted after review: indented blocks (nested functions, methods) are made parseable via the synthetic wrapper while the text itself stays verbatim, so string literals and column offsets are untouched. Explicit source normalization is deferred to arxjit: define a normalized source representation for ExtractedSource / cache keys #97.
• Remove the decorator layer before parsing — see "How stripping works" below
• Parse the function source using Python's built-in ast module
• Provide clear errors when source extraction is unavailable — follow-up PR (already built and stacked on this branch), so each PR stays small. This PR ships the one error case it structurally needs: a retrieved source that is not a single function definition (e.g. a lambda) raises SourceExtractionError with a structured Diagnostic, because returning an ExtractedSource for it would be wrong.

What it returns

ExtractedSource (frozen dataclass):

field	meaning
filename	file defining the (unwrapped) function, or "<unknown>"
source	verbatim file text from the def line to the end of the function; decorator lines removed, original indentation preserved
lineno	1-based line of the def statement in filename
node	parsed ast.FunctionDef | ast.AsyncFunctionDef with an empty decorator list

Two properties worth calling out:

• Node line numbers point into the user's real file (ast.increment_lineno after the parse). The validation pass can put node.lineno straight into a Diagnostic without any offset bookkeeping.
• Column offsets are file-true. Because the text is never dedented, col_offset / end_col_offset index into the real file line (as 0-based UTF-8 byte offsets). Per the column contract from feat(arxjit): add structured diagnostics and the public error hierarchy #95, the only remaining step at the diagnostic-producing boundary is the byte-to-character conversion.

How stripping works

The block returned by inspect.getsourcelines includes the decorators. If its first line is indented, the block is parsed as "if True:\n" + block — this preserves every byte of the original text (multiline strings with zero-indent lines parse fine, unlike with textwrap.dedent). FunctionDef.lineno always points at the def line (decorators live separately in decorator_list), so the verbatim source slice starts there and the node's decorator_list is emptied. The result has no decorators by construction rather than by pattern-matching.

Design decisions to review

1. async def extracts successfully (hence the FunctionDef | AsyncFunctionDef union). Rejecting async is the validation pass's job, so users get a proper "unsupported construct" diagnostic with a location instead of a confusing extraction failure.
2. Diagnostic.code stays None — no stable code scheme (AJ001, ...) exists yet; I'd propose introducing one in the validation PR where the per-construct diagnostics land.
3. extract_source / ExtractedSource are exported from arxjit — useful for the debug/introspection sprint later; easy to make private if you prefer.
4. PyFunc alias is duplicated from core.py instead of imported, to avoid a core <-> source import cycle once the decorator wiring PR makes core import source.

Verification

• 14 extraction tests (47 total), arxjit coverage 100%; line-number assertions read the test file itself so they don't rot when the file is edited. Includes one regression test per review finding: a nested function containing a multiline string with zero-indent lines, file-true col_offset/end_col_offset on a nested function, and extract_source on an actual @jit result.
• The stdlib behaviors the implementation relies on were probed on local CPython 3.10, 3.11, 3.13 and 3.14 (matching the CI matrix): FunctionDef.lineno is the def line under multi-line decorators, increment_lineno shifts end_lineno too, and constructs that are valid syntax but unsupported semantics (nonlocal, yield, methods) parse fine standalone and flow through to validation as intended.
• Verified end-to-end in IPython (the original Colab demo environment): a @jit-decorated function defined in a cell extracts cleanly — decorator stripped, <ipython-input-...> filename (including when extracting the JitFunction itself), cell-relative line numbers.
• All gates green locally: pytest + the 86% coverage gate, ruff format/check, mypy strict, bandit, vulture, mccabe, and douki sync idempotent.

[xmnlab]

thanks @Jaskirat-s7 for working on that

PR review: changes requested

some coments from my side:

1. High — textwrap.dedent breaks valid nested functions containing multiline strings

textwrap.dedent calculates indentation across every physical line, including lines inside triple-quoted strings. A nested function can legally contain string lines at column zero:

def outer():
def nested():
value = """first
second
"""
return value

return nested

For this function, textwrap.dedent removes nothing because second and the closing quotes have no indentation. The subsequent ast.parse() raises IndentationError, even though the original function is valid Python. This directly contradicts the intended support for nested functions and methods.

The indentation should be removed using the function definition’s leading prefix rather than the minimum indentation of the entire block, or the block should initially be parsed inside a synthetic wrapper that preserves lexical content.

Location: packages/arxjit/src/arxjit/source.py:142

2. High — column offsets no longer point into the real source file

Dedenting a method or nested function shifts all AST column offsets to the left. ast.increment_lineno() restores line numbers only; it does not restore col_offset or end_col_offset.

For example:

def outer():
def nested():
return missing

The return begins at zero-based column 8 in the file, but the returned AST reports column 4. Converting that byte offset to a one-based character column still produces the wrong location.

This conflicts with the diagnostic contract, which says diagnostics point into the user’s source and store one-based character columns. The ExtractedSource documentation currently implies the offsets only require byte-to-character conversion, but the removed indentation must also be accounted for.

Location: packages/arxjit/src/arxjit/source.py:142–163

3. Medium — extracting an actual @jit result loses its filename

JitFunction uses functools.update_wrapper, so inspect.getsourcelines() unwraps it and successfully retrieves the original function’s source.

However, _filename_of() runs against the wrapper itself before inspection and only looks for code. JitFunction has no code, so this returns "".

Consequently:

@arxjit.jit
def add(a, b):
return a + b

arxjit.extract_source(add).filename

""

That is especially surprising because the module and public function describe extraction of @jit-decorated functions. Unwrapping once with inspect.unwrap() before deriving the filename would keep source retrieval, naming, and filename attribution consistent.

Location: packages/arxjit/src/arxjit/source.py:140

[Jaskirat-s7]

Thanks for the detailed review @xmnlab! All three findings are addressed in the new commit:

1. textwrap.dedent breaking nested functions with multiline strings -> you're right, and I went with the synthetic-wrapper option you suggested: dedent is gone entirely. If the retrieved block's first line is indented, it's parsed as "if True:\n" + block, so the text is never modified and string literal content is preserved byte-for-byte. I reproduced your exact outer/nested example as a regression test (test_nested_function_with_multiline_string_is_extracted) ,it fails with IndentationError on the old code and passes now, with the parsed constant equal to "first\nsecond\n" verbatim.

2. Column offsets not pointing into the real file : this I fixed as a direct consequence of dropping dedent: since the text keeps its original indentation, col_offset and end_col_offset now index into the real file line with no correction needed. The ExtractedSource docs are updated to say exactly that (only the byte→character conversion remains at the diagnostic boundary), and test_column_offsets_point_into_the_real_file asserts line[col_offset:end_col_offset] == "return x + 1" against the actual file line. A side effect worth noting: ExtractedSource.source is now the verbatim file slice from the def line (original indentation preserved) rather than a dedented copy, so nested-function sources are no longer standalone-parseable , the parsed node is the canonical artifact for the validation pass.

3. JitFunction losing its filename ,this i fixed with your suggestion: _filename_of now does inspect.unwrap(fn) before reading code, matching what getsourcelines does for retrieval, so extract_source(add).filename on an actual @jit result returns the defining file. Regression test test_jit_function_is_extracted_with_filename runs your exact snippet, and I also added a cross-file case (a functools.wraps wrapper compiled under a different filename) to pin the attribution to the wrapped function's file.

14 extraction tests now (one regression test per finding), arxjit coverage still 100%, all gates green locally including douki. Ready for another look!

[xmnlab]

The latest changes address all three previously reported correctness issues, and the added regression tests cover them well. I do not see a code-level blocker to merging this PR.

One contract detail should be handled as a follow-up: "ExtractedSource.source" now preserves the original indentation, so nested-function source is not standalone-parseable and should not currently be described as normalized. Please update the PR description to reflect the new behavior, and create a follow-up issue to either introduce a separate normalized source representation or define explicit normalization when the cache-key work is implemented.

Approved from my side.

[Jaskirat-s7]

Thanks for the approval @xmnlab! Both follow-ups are done:

1.The PR description is updated to reflect the current behavior: ExtractedSource.source is now described as the verbatim file slice (original indentation preserved, not standalone-parseable for nested functions, not a normalized representation), with the parsed node as the canonical artifact for the later stages.
2.i have opened #97 to track the normalized source representation, covering both options you mentioned , a separate normalized form (e.g. ast.unparse of the decorator-free node) or defining explicit normalization as part of the cache-key design when the caching sprint lands. I'd lean toward deciding it there, since the cache key is the first real consumer, but happy to go either way.