⚡️ Speed up function function_kind by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for function_kind in codeflash/code_utils/static_analysis.py

⏱️ Runtime : 1.44 milliseconds → 1.30 milliseconds (best of 160 runs)

📝 Explanation and details

Here is an optimized version of your function, making the following performance improvements.

• Avoids repeatedly creating lists and checking list membership (parents[0].type in ["FunctionDef", "AsyncFunctionDef"]). This is more efficiently done via a set: {"FunctionDef", "AsyncFunctionDef"}.
• Reduces attribute lookups and loop overhead inside the decorator check, using local variables, early return, and prepares the valid decorator names as a set for speedy comparison.
• Reduces repeated attribute lookups (parents[0]), storing it in a local variable for speed.
• Avoids unnecessary isinstance() checks/loops by pre-filtering.
• Keeps both return type and functionality the same as before.

Changes Made

• Changed ["FunctionDef", "AsyncFunctionDef"] to {"FunctionDef", "AsyncFunctionDef"} for O(1) lookup.
• Reduced multiple attribute lookups by storing parents[0] in parent_type.
• Combined the if and elif for decorator names for faster path to result (early return).
• The logic and return values are identical.

This should result in reduced runtime, particularly for the high-frequency paths in your line profiling results.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 5542 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests Details

from __future__ import annotations

import ast
# Enum for function kind
from enum import Enum, auto
from typing import Any

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.static_analysis import function_kind

# --- Mocked/Minimal dependencies for testability ---


class FunctionKind(Enum):
    FUNCTION = auto()
    INSTANCE_METHOD = auto()
    CLASS_METHOD = auto()
    STATIC_METHOD = auto()

# Minimal FunctionParent class to mimic the real one
class FunctionParent:
    def __init__(self, type_: str):
        self.type = type_
from codeflash.code_utils.static_analysis import function_kind

# --- Unit tests ---

# Helper to build a FunctionDef/AsyncFunctionDef with decorators
def make_func(name="f", decorators=None, async_=False):
    if decorators is None:
        decorators = []
    node_type = ast.AsyncFunctionDef if async_ else ast.FunctionDef
    return node_type(
        name=name,
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=[],
        decorator_list=decorators,
        returns=None,
        type_comment=None
    )

# --------------------------
# 1. Basic Test Cases
# --------------------------

def test_top_level_function():
    """Test a top-level function (no parents)."""
    node = make_func()
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_top_level_async_function():
    """Test a top-level async function (no parents)."""
    node = make_func(async_=True)
    parents = []
    codeflash_output = function_kind(node, parents) # 334ns -> 375ns

def test_nested_function():
    """Test a function nested inside another function."""
    node = make_func()
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 541ns

def test_nested_async_function():
    """Test an async function nested inside another async function."""
    node = make_func(async_=True)
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_instance_method():
    """Test a function defined inside a class (no decorators)."""
    node = make_func()
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 625ns -> 667ns

def test_class_method():
    """Test a @classmethod inside a class."""
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 791ns -> 791ns

def test_static_method():
    """Test a @staticmethod inside a class."""
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 833ns -> 833ns

def test_multiple_decorators_staticmethod_last():
    """Test multiple decorators, @staticmethod is last."""
    decorator1 = ast.Name(id="foo", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 959ns -> 1.00μs

def test_multiple_decorators_classmethod_first():
    """Test multiple decorators, @classmethod is first."""
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="bar", ctx=ast.Load())
    node = make_func(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 792ns -> 791ns

# --------------------------
# 2. Edge Test Cases
# --------------------------

def test_no_parents_list():
    """Test with parents=None (should treat as top-level function)."""
    node = make_func()
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_unknown_parent_type():
    """Test with a parent type that is not recognized."""
    node = make_func()
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 417ns -> 417ns

def test_empty_decorator_list_in_class():
    """Test function in class with empty decorator_list."""
    node = make_func(decorators=[])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 625ns -> 666ns

def test_non_name_decorator_in_class():
    """Test function in class with a decorator that's not ast.Name."""
    decorator = ast.Attribute(value=ast.Name(id="foo", ctx=ast.Load()), attr="bar", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    # Should still be instance method since no class/static method decorator found
    codeflash_output = function_kind(node, parents) # 792ns -> 833ns

def test_both_class_and_static_method_decorators():
    """Test function with both @classmethod and @staticmethod (should prefer first found)."""
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 792ns -> 750ns

def test_both_static_and_class_method_decorators():
    """Test function with both @staticmethod and @classmethod (should prefer first found)."""
    decorator1 = ast.Name(id="staticmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 833ns -> 791ns

def test_function_with_weird_decorator_names():
    """Test function with decorators that are not classmethod/staticmethod."""
    decorator1 = ast.Name(id="property", ctx=ast.Load())
    decorator2 = ast.Name(id="custom", ctx=ast.Load())
    node = make_func(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 917ns -> 1.00μs


def test_function_with_empty_parents_but_module_parent():
    """Test function with parents list containing a Module parent."""
    node = make_func()
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 458ns -> 458ns

def test_function_with_multiple_parents_first_is_function():
    """Test function with multiple parents, first is function."""
    node = make_func()
    parents = [FunctionParent("FunctionDef"), FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 459ns -> 583ns

def test_function_with_multiple_parents_first_is_class():
    """Test function with multiple parents, first is class."""
    node = make_func()
    parents = [FunctionParent("ClassDef"), FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents) # 625ns -> 667ns

def test_function_with_multiple_parents_first_is_unknown():
    """Test function with multiple parents, first is unknown."""
    node = make_func()
    parents = [FunctionParent("UnknownType"), FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 459ns -> 417ns

# --------------------------
# 3. Large Scale Test Cases
# --------------------------

def test_large_batch_of_functions():
    """Test a large number of top-level functions (performance and correctness)."""
    nodes = [make_func(name=f"f{i}") for i in range(500)]
    parents = []
    for node in nodes:
        codeflash_output = function_kind(node, parents) # 83ns -> 83ns

def test_large_batch_of_class_methods():
    """Test a large number of class methods in a class."""
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    nodes = [make_func(name=f"cm{i}", decorators=[decorator]) for i in range(500)]
    parents = [FunctionParent("ClassDef")]
    for node in nodes:
        codeflash_output = function_kind(node, parents) # 250ns -> 250ns

def test_large_batch_of_static_methods():
    """Test a large number of static methods in a class."""
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    nodes = [make_func(name=f"sm{i}", decorators=[decorator]) for i in range(500)]
    parents = [FunctionParent("ClassDef")]
    for node in nodes:
        codeflash_output = function_kind(node, parents) # 291ns -> 250ns

def test_large_batch_of_instance_methods_with_varied_decorators():
    """Test a large number of instance methods with varied (non-matching) decorators."""
    # Use decorators that are not classmethod or staticmethod
    decorators = [ast.Name(id=f"decorator{i}", ctx=ast.Load()) for i in range(5)]
    parents = [FunctionParent("ClassDef")]
    for i in range(500):
        node = make_func(name=f"im{i}", decorators=[decorators[i % 5]])
        codeflash_output = function_kind(node, parents) # 250ns -> 208ns

def test_large_batch_of_nested_functions():
    """Test a large batch of nested functions (function inside function)."""
    parents = [FunctionParent("FunctionDef")]
    for i in range(500):
        node = make_func(name=f"nf{i}")
        codeflash_output = function_kind(node, parents) # 125ns -> 125ns
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import ast
from enum import Enum
from typing import List

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.static_analysis import function_kind


# Simulate the FunctionKind enum as expected by the function
class FunctionKind(Enum):
    FUNCTION = "function"
    CLASS_METHOD = "class_method"
    STATIC_METHOD = "static_method"
    INSTANCE_METHOD = "instance_method"

# Simulate the FunctionParent dataclass as expected by the function
class FunctionParent:
    def __init__(self, type_: str):
        self.type = type_
from codeflash.code_utils.static_analysis import function_kind

# --------- Unit Tests ---------

# Helper to create a FunctionDef node with decorators
def make_funcdef(name, decorators=None, async_=False):
    if decorators is None:
        decorators = []
    if async_:
        return ast.AsyncFunctionDef(
            name=name,
            args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )
    else:
        return ast.FunctionDef(
            name=name,
            args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )

# --------- Basic Test Cases ---------

def test_top_level_function():
    # Test: top-level function (no parents)
    node = make_funcdef("foo")
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_top_level_async_function():
    # Test: top-level async function (no parents)
    node = make_funcdef("foo", async_=True)
    parents = []
    codeflash_output = function_kind(node, parents) # 334ns -> 375ns

def test_nested_function():
    # Test: function nested within another function
    node = make_funcdef("bar")
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 541ns

def test_nested_async_function():
    # Test: async function nested within another async function
    node = make_funcdef("bar", async_=True)
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_instance_method():
    # Test: method inside a class, no decorators
    node = make_funcdef("baz")
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 625ns -> 667ns

def test_class_method():
    # Test: method inside a class, with @classmethod decorator
    decorator = ast.Name(id="classmethod")
    node = make_funcdef("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 791ns -> 791ns

def test_static_method():
    # Test: method inside a class, with @staticmethod decorator
    decorator = ast.Name(id="staticmethod")
    node = make_funcdef("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 833ns -> 833ns

# --------- Edge Test Cases ---------

def test_multiple_decorators_classmethod_first():
    # Test: method with multiple decorators, @classmethod is first
    decorators = [ast.Name(id="classmethod"), ast.Name(id="otherdecorator")]
    node = make_funcdef("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 792ns -> 791ns

def test_multiple_decorators_staticmethod_second():
    # Test: method with multiple decorators, @staticmethod is second
    decorators = [ast.Name(id="other"), ast.Name(id="staticmethod")]
    node = make_funcdef("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.04μs -> 1.00μs

def test_decorator_with_attribute():
    # Test: method with decorator as ast.Attribute (should be ignored)
    decorator = ast.Attribute(value=ast.Name(id="foo"), attr="classmethod")
    node = make_funcdef("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    # Should fall back to instance method, since only ast.Name is considered
    codeflash_output = function_kind(node, parents) # 750ns -> 833ns

def test_decorator_with_call():
    # Test: method with decorator as ast.Call (should be ignored)
    decorator = ast.Call(func=ast.Name(id="classmethod"), args=[], keywords=[])
    node = make_funcdef("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    # Should fall back to instance method, since only ast.Name is considered
    codeflash_output = function_kind(node, parents) # 750ns -> 791ns

def test_unexpected_parent_type():
    # Test: parent type that is not handled (should return None)
    node = make_funcdef("baz")
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 417ns -> 416ns

def test_empty_parents_list():
    # Test: empty parents list (should be function)
    node = make_funcdef("foo")
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_parent_not_first_classdef():
    # Test: parent list where first is not ClassDef, but later is
    node = make_funcdef("baz")
    parents = [FunctionParent("FunctionDef"), FunctionParent("ClassDef")]
    # Only the first parent matters; should be FUNCTION
    codeflash_output = function_kind(node, parents) # 500ns -> 542ns

def test_unknown_decorator_name():
    # Test: method with unknown decorator, should be instance method
    decorator = ast.Name(id="unknown_decorator")
    node = make_funcdef("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 833ns -> 833ns

def test_async_method_in_class():
    # Test: async method inside a class, no decorators
    node = make_funcdef("baz", async_=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 625ns -> 625ns

def test_async_class_method():
    # Test: async method inside a class, with @classmethod decorator
    decorator = ast.Name(id="classmethod")
    node = make_funcdef("baz", decorators=[decorator], async_=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 791ns -> 791ns

def test_async_static_method():
    # Test: async method inside a class, with @staticmethod decorator
    decorator = ast.Name(id="staticmethod")
    node = make_funcdef("baz", decorators=[decorator], async_=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 875ns -> 792ns

# --------- Large Scale Test Cases ---------

def test_many_decorators_with_classmethod_at_end():
    # Test: method with many decorators, @classmethod at the end
    decorators = [ast.Name(id=f"decorator_{i}") for i in range(998)]
    decorators.append(ast.Name(id="classmethod"))
    node = make_funcdef("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 56.9μs -> 49.7μs

def test_many_decorators_with_staticmethod_in_middle():
    # Test: method with many decorators, @staticmethod in the middle
    decorators = [ast.Name(id=f"decorator_{i}") for i in range(500)]
    decorators.insert(250, ast.Name(id="staticmethod"))
    node = make_funcdef("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 15.8μs -> 14.2μs

def test_large_parents_list_first_is_classdef():
    # Test: long parents list, first is ClassDef
    node = make_funcdef("baz")
    parents = [FunctionParent("ClassDef")] + [FunctionParent("FunctionDef")] * 999
    codeflash_output = function_kind(node, parents) # 625ns -> 667ns

def test_large_parents_list_first_is_functiondef():
    # Test: long parents list, first is FunctionDef
    node = make_funcdef("baz")
    parents = [FunctionParent("FunctionDef")] + [FunctionParent("ClassDef")] * 999
    codeflash_output = function_kind(node, parents) # 458ns -> 542ns

def test_large_number_of_functions():
    # Test: many top-level functions in a loop
    for i in range(1000):
        node = make_funcdef(f"func_{i}")
        parents = []
        codeflash_output = function_kind(node, parents) # 83ns -> 83ns

def test_large_number_of_class_methods():
    # Test: many class methods in a class
    for i in range(1000):
        decorator = ast.Name(id="classmethod")
        node = make_funcdef(f"func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents) # 250ns -> 209ns

def test_large_number_of_static_methods():
    # Test: many static methods in a class
    for i in range(1000):
        decorator = ast.Name(id="staticmethod")
        node = make_funcdef(f"func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents) # 250ns -> 208ns
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-function_kind-mbslqf6c and push.