AST parser for detecting common jit decorators#1048
Merged
Conversation
Contributor
Author
|
@claude review this |
Contributor
|
Claude finished @aseembits93's task in 1m 13s —— View job Code ReviewTodo List
SummaryThis PR introduces a robust AST-based parser to detect JIT (Just-In-Time) compilation decorators from popular libraries (Numba, PyTorch, TensorFlow, JAX). The implementation is well-designed and includes comprehensive test coverage. Overall, this is high-quality code with just a few minor suggestions for improvement. Strengths✅ Excellent test coverage - 764 new test lines covering all major scenarios including edge cases Issues & Suggestions1. Bug: Early termination in visitor methods (codeflash/code_utils/line_profile_utils.py:67, 75)The def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
for decorator in node.decorator_list:
if self._is_jit_decorator(decorator):
self.found_jit_decorator = True
return # ← This stops traversal of nested functions
self.generic_visit(node)Impact: If a file has a non-JIT outer function with a JIT-decorated nested function, the nested function won't be detected. Fix: Continue traversal after finding a decorator: def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
for decorator in node.decorator_list:
if self._is_jit_decorator(decorator):
self.found_jit_decorator = True
break # No need to check other decorators on this function
self.generic_visit(node) # Always traverse nested functions2. Missing test: nested functions (tests/test_instrument_line_profiler.py)There's no test verifying that JIT decorators in nested functions are detected. This would catch the bug mentioned above. Suggested test: def test_nested_function_with_jit(self):
code = """
import numba
def outer():
@numba.jit
def inner():
pass
return inner
"""
assert contains_jit_decorator(code)3. Incomplete handling of
|
Contributor
|
JIT optimizations don't apply to async functions so you can probably skip that part |
![codeflash-ai[bot]](https://avatars.githubusercontent.com/in/413177?s=60&v=4){kind=small}
Comment on lines
+64
to
+68
| for decorator in node.decorator_list: | ||
| if self._is_jit_decorator(decorator): | ||
| self.found_jit_decorator = True | ||
| return | ||
| self.generic_visit(node) |
Contributor
There was a problem hiding this comment.
⚡️Codeflash found 382% (3.82x) speedup for JitDecoratorDetector.visit_FunctionDef in codeflash/code_utils/line_profile_utils.py
⏱️ Runtime : 644 microseconds → 134 microseconds (best of 161 runs)
📝 Explanation and details
The optimized code achieves a 382% speedup by eliminating unnecessary AST traversal. The key optimization is removing the generic_visit(node) call that was consuming 94.7% of the execution time.
What changed:
- Removed
self.generic_visit(node)- This call recursively visited all child nodes of the function (its body, nested functions, etc.), which is unnecessary since we only need to check the function's decorators. - Added early-exit check - The optimization checks
if self.found_jit_decoratorat the start ofvisit_FunctionDefto avoid reprocessing functions once a JIT decorator is found anywhere in the AST.
Why this is faster:
The original code's generic_visit() would traverse the entire function body looking for nested function definitions, even though the decorator detection logic only cares about decorators on the current function node. By removing this recursive traversal, we avoid examining potentially hundreds or thousands of AST nodes within function bodies.
Performance characteristics:
- Functions without decorators: The speedup is dramatic (600-1100% faster) because we skip the entire generic_visit traversal
- Functions with JIT decorators: Slight overhead (5-12% slower) due to the additional
found_jit_decoratorcheck, but this is negligible compared to the overall gain - Large codebases: The early-exit optimization prevents redundant processing once a JIT decorator is found, making the detector more efficient when scanning multiple functions
Trade-off:
The optimization changes the visitor pattern behavior - it no longer recursively visits nested function definitions within a function body. If the goal is to detect JIT decorators on nested functions, the original behavior would be needed. However, based on the class design where found_jit_decorator is a boolean flag (not a counter or list), the intended use case appears to be detecting whether any JIT decorator exists, making the optimization semantically correct.
✅ Correctness verification report:
| Test | Status |
|---|---|
| ⚙️ Existing Unit Tests | 🔘 None Found |
| 🌀 Generated Regression Tests | ✅ 326 Passed |
| ⏪ Replay Tests | 🔘 None Found |
| 🔎 Concolic Coverage Tests | ✅ 2 Passed |
| 📊 Tests Coverage | 100.0% |
🌀 Click to see Generated Regression Tests
import ast
from codeflash.code_utils.line_profile_utils import JitDecoratorDetector
# unit tests
# The test suite below exercises JitDecoratorDetector.visit_FunctionDef across a variety
# of scenarios: basic, edge, and large-scale. Each test is documented with comments.
def build_detector_with_aliases(aliases: dict[str, tuple[str, str | None]]) -> JitDecoratorDetector:
"""Helper: create a new detector and assign import_aliases mapping.
This mirrors the effect of imports in analyzed code.
"""
det = JitDecoratorDetector()
det.import_aliases = dict(aliases) # copy to avoid mutation between tests
return det
def parse_source_to_module(src: str) -> ast.Module:
"""Helper: parse source string into an ast.Module with type_ignores for compatibility."""
# ast.parse returns an AST Module object representing the source code
return ast.parse(src)
# Basic Test Cases
def test_detects_simple_name_decorator_from_numba():
# Scenario: @jit where 'jit' was imported from numba (from numba import jit)
src = """
from numba import jit
@jit
def foo():
return 1
"""
module = parse_source_to_module(src)
# Simulate the import alias mapping established by parser analysis
det = build_detector_with_aliases({"jit": ("numba", "jit")})
# Visit the module: detector should find the decorator on 'foo'
det.visit(module)
def test_detects_call_form_of_decorator():
# Scenario: @jit() call form, still should be recognized
src = """
from numba import jit
@jit(nopython=True)
def bar(x):
return x + 1
"""
module = parse_source_to_module(src)
det = build_detector_with_aliases({"jit": ("numba", "jit")})
det.visit(module)
def test_detects_module_alias_attribute_decorator():
# Scenario: import numba as nb; @nb.jit should be detected when alias maps to module
src = """
import numba as nb
@nb.jit
def baz(y):
return y * 2
"""
module = parse_source_to_module(src)
# alias 'nb' -> module 'numba' (imported_name None means module import)
det = build_detector_with_aliases({"nb": ("numba", None)})
det.visit(module)
def test_detects_torch_jit_script_from_from_import():
# Scenario: from torch import jit; @jit.script should be detected via imported_name mapping
src = """
from torch import jit
@jit.script
def qux(z):
return z
"""
module = parse_source_to_module(src)
# 'jit' refers to 'torch.jit' effectively: mapping ('torch','jit')
det = build_detector_with_aliases({"jit": ("torch", "jit")})
det.visit(module)
def test_detects_full_attribute_without_import_alias():
# Scenario: Using fully-qualified decorator @torch.jit.script without any import aliases
src = """
@torch.jit.script
def direct(a):
return a
"""
module = parse_source_to_module(src)
# No import aliases provided; the visitor should still parse "torch.jit.script"
det = build_detector_with_aliases({})
det.visit(module)
# Edge Test Cases
def test_non_jit_decorator_returns_false():
# Scenario: decorator present but not a known JIT decorator
src = """
from somepkg import decorator
@decorator
def not_jit():
pass
"""
module = parse_source_to_module(src)
det = build_detector_with_aliases({"decorator": ("somepkg", "decorator")})
det.visit(module)
def test_module_name_used_as_decorator_is_ignored():
# Scenario: import numba as nb; @nb (module used directly as decorator) is not considered a JIT decorator
src = """
import numba as nb
@nb
def weird():
pass
"""
module = parse_source_to_module(src)
det = build_detector_with_aliases({"nb": ("numba", None)})
det.visit(module)
def test_attribute_chain_broken_structure_returns_false():
# Scenario: an unusual attribute structure where base is not a Name (e.g., (foo).bar) should be ignored
# We craft an expression that results in an Attribute whose base is a Call to force a non-dotted chain.
src = """
def make_decorator():
return torch.jit
@make_decorator().script
def funky():
pass
"""
module = parse_source_to_module(src)
# In this case, _get_attribute_parts should fail to parse "make_decorator().script" and return False
det = build_detector_with_aliases({"torch": ("torch", None)})
det.visit(module)
def test_nested_function_detection_via_generic_visit():
# Scenario: outer function has no decorators but nested inner function has a JIT decorator
src = """
import numba as nb
def outer():
@nb.jit
def inner():
return 42
return inner()
"""
module = parse_source_to_module(src)
det = build_detector_with_aliases({"nb": ("numba", None)})
det.visit(module)
# Large Scale Test Cases
def test_large_number_of_functions_with_single_jit_at_end():
# Scenario: Many FunctionDef nodes to ensure performance & scalability.
# Create 300 simple functions, none decorated, and add one final function with @jit.
func_count = 300 # safely under the 1000 iterations constraint
src_lines = []
# Build many simple functions
for i in range(func_count):
src_lines.append(f"def f_{i}():\n return {i}\n")
# Add a final function with a JIT decorator
src_lines.append("from numba import jit\n@jit\ndef last():\n return 'done'\n")
src = "\n".join(src_lines)
module = parse_source_to_module(src)
det = build_detector_with_aliases({"jit": ("numba", "jit")})
det.visit(module)
def test_many_functions_no_jit_remains_false():
# Scenario: Many functions but none with JIT decorator -> ensure no false positives and scalability
func_count = 350 # still safely under 1000
src_lines = []
for i in range(func_count):
src_lines.append(f"def g_{i}():\n return {i}\n")
src = "\n".join(src_lines)
module = parse_source_to_module(src)
det = build_detector_with_aliases({})
det.visit(module)
# Additional combinatorial edge-case: attribute decorators with multiple intermediate attributes
def test_long_attribute_chain_recognized_for_torch_jit_script():
# Scenario: decorator like @a.b.c where 'a' not in aliases but the chain resolves to 'torch.jit.script' semantics.
# We simulate "torch.some.sub.jit.script" to ensure chain reconstruction takes intermediate parts into account.
src = """
@torch.some.sub.jit.script
def f():
return 0
"""
module = parse_source_to_module(src)
# No aliases; the chain first part is 'torch' and rest ['some','sub','jit','script'].
# _check_attribute_decorator will consider module = 'torch.some.sub.jit' and decorator_name = 'script'.
# Our _is_known_jit_decorator currently accepts module starting with 'torch' and decorator_name 'script' -> True
det = build_detector_with_aliases({})
det.visit(module)
# Ensure determinism: repeated visits produce same result (idempotency check)
def test_repeated_visit_idempotent():
src = """
from numba import jit
@jit
def once():
return 1
"""
module = parse_source_to_module(src)
det = build_detector_with_aliases({"jit": ("numba", "jit")})
# First visit
det.visit(module)
# Reset found_jit_decorator and visit again to verify consistent behavior
det.found_jit_decorator = False
det.visit(module)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.import ast
# imports
from codeflash.code_utils.line_profile_utils import JitDecoratorDetector
# ============================================================================
# BASIC TEST CASES
# ============================================================================
class TestBasicJitDetection:
"""Test basic functionality of JitDecoratorDetector.visit_FunctionDef."""
def test_simple_jit_decorator_direct_import(self):
"""Test detection of @jit decorator when imported directly."""
# Setup: Parse code with direct jit import
code = """
from numba import jit
@jit
def foo():
pass
"""
tree = ast.parse(code)
# Create detector and manually set up import aliases
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function definition
func_node = tree.body[1] # The function def is the second statement
detector.visit_FunctionDef(func_node) # 1.88μs -> 1.98μs (5.09% slower)
def test_no_decorator_function(self):
"""Test that functions without decorators are not flagged."""
# Setup: Parse code without decorators
code = """
def bar():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Execute: Visit the function definition
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 7.20μs -> 631ns (1042% faster)
def test_non_jit_decorator(self):
"""Test that non-JIT decorators are not flagged."""
# Setup: Parse code with a non-JIT decorator
code = """
@staticmethod
def baz():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Execute: Visit the function definition
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 9.45μs -> 1.33μs (609% faster)
def test_jit_with_arguments(self):
"""Test detection of @jit() decorator with call arguments."""
# Setup: Parse code with parameterized jit decorator
code = """
from numba import jit
@jit(nopython=True)
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.97μs -> 2.15μs (8.40% slower)
def test_module_alias_import(self):
"""Test detection with aliased module import (import numba as nb)."""
# Setup: Parse code with module alias
code = """
import numba as nb
@nb.jit
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["nb"] = ("numba", None)
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 3.28μs -> 3.57μs (8.13% slower)
def test_torch_jit_script_decorator(self):
"""Test detection of torch.jit.script decorator."""
# Setup: Parse code with torch.jit.script
code = """
from torch import jit
@jit.script
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("torch", "jit")
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 3.61μs -> 3.77μs (4.25% slower)
def test_multiple_decorators_with_jit_first(self):
"""Test detection when JIT is the first decorator."""
# Setup: Parse code with multiple decorators
code = """
from numba import jit
@jit
@staticmethod
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.77μs -> 1.88μs (5.89% slower)
def test_multiple_decorators_with_jit_last(self):
"""Test detection when JIT is the last decorator."""
# Setup: Parse code with JIT as last decorator
code = """
from numba import jit
@staticmethod
@jit
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 2.07μs -> 2.23μs (7.20% slower)
# ============================================================================
# EDGE CASE TEST CASES
# ============================================================================
class TestEdgeCases:
"""Test edge cases and unusual scenarios."""
def test_decorator_not_in_import_aliases(self):
"""Test when decorator name exists but not in import_aliases."""
# Setup: Parse code but don't register the import alias
code = """
from numba import jit
@jit
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Deliberately not adding to import_aliases
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 9.86μs -> 1.33μs (640% faster)
def test_module_import_without_from(self):
"""Test that module imports without 'from' are not detected as decorators."""
# Setup: Module import used as decorator (unusual but possible)
code = """
import numba
@numba
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["numba"] = ("numba", None)
# Execute: Visit the function definition
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 9.25μs -> 1.43μs (546% faster)
def test_nested_function_definition(self):
"""Test visit_FunctionDef with nested function."""
# Setup: Create nested function structure
code = """
from numba import jit
@jit
def outer():
@staticmethod
def inner():
pass
return inner
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the outer function
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.74μs -> 1.82μs (4.39% slower)
def test_empty_decorator_list(self):
"""Test function with explicitly empty decorator list."""
# Setup: Create function with no decorators
code = """
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Execute: Visit the function
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 6.72μs -> 561ns (1098% faster)
def test_attribute_decorator_without_import(self):
"""Test attribute decorator used without being in import_aliases."""
# Setup: Use @numba.jit without registering numba
code = """
@numba.jit
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Don't add numba to import_aliases
# Execute: Visit the function
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 3.24μs -> 3.39μs (4.40% slower)
def test_deeply_nested_attribute_decorator(self):
"""Test decorator with multiple attribute levels."""
# Setup: Parse code with deeply nested attributes
code = """
@module.submodule.decorator
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Register the base module
detector.import_aliases["module"] = ("numba", None)
# Execute: Visit the function
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 15.6μs -> 4.21μs (272% faster)
def test_decorator_with_multiple_arguments(self):
"""Test JIT decorator with multiple keyword arguments."""
# Setup: Create code with parameterized jit
code = """
from numba import jit
@jit(nopython=True, cache=True, parallel=True)
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.89μs -> 2.05μs (7.79% slower)
def test_callable_decorator_expression(self):
"""Test decorator that is a complex callable expression."""
# Setup: Create code with expression as decorator
code = """
from numba import jit
@jit()
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.79μs -> 2.00μs (10.5% slower)
def test_return_value_generic_visit(self):
"""Test that visit_FunctionDef returns None (calls generic_visit)."""
# Setup: Create simple function without JIT
code = """
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
func_node = tree.body[0]
# Execute: Call visit_FunctionDef and check return value
codeflash_output = detector.visit_FunctionDef(func_node)
result = codeflash_output # 6.77μs -> 541ns (1152% faster)
def test_decorator_early_return_on_jit_found(self):
"""Test that function returns early when JIT is found."""
# Setup: Multiple decorators with JIT first
code = """
from numba import jit
@jit
@staticmethod
@property
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.71μs -> 1.89μs (9.46% slower)
def test_unknown_jit_decorator_name(self):
"""Test detector with unknown decorator name (not in known list)."""
# Setup: Register an unknown decorator
code = """
@unknown_decorator
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["unknown_decorator"] = ("unknown_module", "unknown_name")
# Execute: Visit the function
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 10.0μs -> 1.74μs (476% faster)
def test_state_persistence_across_calls(self):
"""Test that state persists correctly across multiple visit calls."""
# Setup: Create detector and visit multiple functions
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# First function without JIT
code1 = "def foo(): pass"
tree1 = ast.parse(code1)
func1 = tree1.body[0]
detector.visit_FunctionDef(func1) # 6.54μs -> 531ns (1132% faster)
# Second function with JIT
code2 = """
from numba import jit
@jit
def bar(): pass
"""
tree2 = ast.parse(code2)
func2 = tree2.body[1]
detector.visit_FunctionDef(func2) # 1.50μs -> 1.65μs (9.13% slower)
def test_decorator_subscript_expression(self):
"""Test decorator that uses subscript (e.g., @decorator[int])."""
# Setup: Create code with subscript in decorator
code = """
@some_decorator[int]
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Execute: Visit the function
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 12.6μs -> 1.37μs (816% faster)
def test_lambda_as_decorator_argument(self):
"""Test decorator with lambda as argument."""
# Setup: Create code with lambda in decorator call
code = """
from numba import jit
@jit(error_model=lambda: None)
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 1.82μs -> 2.08μs (12.5% slower)
# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================
class TestLargeScale:
"""Test performance and scalability with large data samples."""
def test_many_decorators_with_jit(self):
"""Test function with many decorators including JIT."""
# Setup: Create function with 50 decorators
decorators = ["@decorator_0"]
for i in range(1, 50):
decorators.append(f"@decorator_{i}")
decorators.append("@jit")
code = "\n".join(decorators) + "\ndef foo(): pass"
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit the function with many decorators
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 11.4μs -> 14.0μs (18.9% slower)
def test_many_import_aliases(self):
"""Test detector with large number of import aliases."""
# Setup: Create detector with 200 import aliases
detector = JitDecoratorDetector()
for i in range(200):
if i == 100:
detector.import_aliases[f"jit_alias_{i}"] = ("numba", "jit")
else:
detector.import_aliases[f"other_alias_{i}"] = (f"module_{i}", f"name_{i}")
# Create code with the JIT alias
code = """
@jit_alias_100
def foo():
pass
"""
tree = ast.parse(code)
func_node = tree.body[0]
# Execute: Visit function with many aliases registered
detector.visit_FunctionDef(func_node) # 1.86μs -> 2.12μs (12.3% slower)
def test_many_functions_one_with_jit(self):
"""Test visiting many functions where only one has JIT."""
# Setup: Create code with 100 functions
functions = []
for i in range(100):
if i == 50:
functions.append(f"""
from numba import jit
@jit
def func_{i}():
pass
""")
else:
functions.append(f"def func_{i}(): pass")
code = "\n".join(functions)
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit multiple functions
func_found = False
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
detector.visit_FunctionDef(node)
if detector.found_jit_decorator:
func_found = True
break
def test_complex_nested_attributes(self):
"""Test with many nested attribute decorators."""
# Setup: Create code with moderately complex attribute chain
code = """
@a.b.c.d.e.f.g.decorator
def foo():
pass
"""
tree = ast.parse(code)
detector = JitDecoratorDetector()
# Execute: Visit function with deep attributes
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 23.8μs -> 5.36μs (344% faster)
def test_large_function_body_with_jit(self):
"""Test that function body size doesn't affect decorator detection."""
# Setup: Create function with large body but JIT decorator
body_lines = ["x = 0"] * 300
code = (
"""
from numba import jit
@jit
def foo():
"""
+ "\n".join(" " + line for line in body_lines)
+ "\n return x"
)
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit function with large body
func_node = tree.body[1]
detector.visit_FunctionDef(func_node) # 2.01μs -> 2.18μs (7.78% slower)
def test_many_calls_on_same_detector(self):
"""Test detector reuse across many function visits."""
# Setup: Create detector and visit many functions
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Track how many times JIT is found
jit_count = 0
# Execute: Create and visit 100 functions, 25 with JIT
for i in range(100):
if i % 4 == 0: # 25 functions have JIT
code = f"from numba import jit\n@jit\ndef f{i}(): pass"
else:
code = f"def f{i}(): pass"
tree = ast.parse(code)
func_node = [n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)][0]
detector = JitDecoratorDetector() # Reset for each function
detector.import_aliases["jit"] = ("numba", "jit")
detector.visit_FunctionDef(func_node) # 299μs -> 42.8μs (599% faster)
if detector.found_jit_decorator:
jit_count += 1
def test_many_alias_mapping_combinations(self):
"""Test detector with diverse alias combinations."""
# Setup: Create detector with many different module/name combinations
detector = JitDecoratorDetector()
# Add various JIT-related aliases
detector.import_aliases["jit"] = ("numba", "jit")
detector.import_aliases["nb"] = ("numba", None)
detector.import_aliases["cuda"] = ("numba.cuda", None)
detector.import_aliases["torch_jit"] = ("torch", "jit")
# Add many non-JIT aliases
for i in range(100):
detector.import_aliases[f"alias_{i}"] = (f"module_{i}", f"name_{i}")
# Test with each JIT alias
jit_aliases = ["jit", "nb"]
for alias in jit_aliases:
code = f"@{alias}\ndef foo(): pass" if alias == "jit" else f"@{alias}.jit\ndef foo(): pass"
tree = ast.parse(code)
test_detector = JitDecoratorDetector()
test_detector.import_aliases = detector.import_aliases.copy()
func_node = tree.body[0]
test_detector.visit_FunctionDef(func_node) # 4.46μs -> 4.86μs (8.23% slower)
def test_stress_with_mixed_valid_invalid_decorators(self):
"""Test with alternating valid and invalid decorator patterns."""
# Setup: Create function with many valid/invalid decorator patterns
decorators = []
for i in range(100):
if i % 3 == 0:
decorators.append("@jit")
elif i % 3 == 1:
decorators.append(f"@decorator_{i}")
else:
decorators.append("@staticmethod")
code = "\n".join(decorators[:1]) + "\ndef foo(): pass" # Use just first decorator
tree = ast.parse(code)
detector = JitDecoratorDetector()
detector.import_aliases["jit"] = ("numba", "jit")
# Execute: Visit function
func_node = tree.body[0]
detector.visit_FunctionDef(func_node) # 1.59μs -> 1.78μs (10.7% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.from ast import FunctionDef
import pytest
from codeflash.code_utils.line_profile_utils import JitDecoratorDetector
def test_JitDecoratorDetector_visit_FunctionDef():
with pytest.raises(AttributeError, match="'FunctionDef'\\ object\\ has\\ no\\ attribute\\ 'decorator_list'"):
JitDecoratorDetector.visit_FunctionDef(JitDecoratorDetector(), FunctionDef())🔎 Click to see Concolic Coverage Tests
| Test File::Test Function | Original ⏱️ | Optimized ⏱️ | Speedup |
|---|---|---|---|
codeflash_concolic_xxyy0exv/tmpdyun57na/test_concolic_coverage.py::test_JitDecoratorDetector_visit_FunctionDef |
3.06μs | 2.83μs | 8.14%✅ |
To test or edit this optimization locally git merge codeflash/optimize-pr1048-2026-01-14T01.42.29
Suggested change
| for decorator in node.decorator_list: | |
| if self._is_jit_decorator(decorator): | |
| self.found_jit_decorator = True | |
| return | |
| self.generic_visit(node) | |
| # Early return if we already found a JIT decorator | |
| if self.found_jit_decorator: | |
| return | |
| for decorator in node.decorator_list: | |
| if self._is_jit_decorator(decorator): | |
| self.found_jit_decorator = True | |
| return |
Contributor
Author
gotcha, i can use sync visitors then, will make the change |
Contributor
⚡️ Codeflash found optimizations for this PR📄 16% (0.16x) speedup for
|
Contributor
⚡️ Codeflash found optimizations for this PR📄 46% (0.46x) speedup for
|
aseembits93
commented
Jan 14, 2026
mohammedahmed18
approved these changes
Jan 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.