⚡️ Speed up function function_kind by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for function_kind in codeflash/code_utils/static_analysis.py

⏱️ Runtime : 494 microseconds → 468 microseconds (best of 69 runs)

📝 Explanation and details

Here’s an optimized rewrite focusing on.

• Avoiding repeated list/dict construction: Replaces ["FunctionDef", "AsyncFunctionDef"] with a set and a local variable for quick lookup.
• Decorator matching: Replaces the for-loop with a single pass using generator expression – thus returning immediately when a match is found, and avoiding extra checks.
• Parent type check: Caches parents[0].type in a variable.
• Reduces function calls: Pulls common checks out of loops.

Here's the optimized code (existing comments are unchanged, since logic is nearly the same).

Summary of improvements.

• Eliminate duplicated string list creation per call.
• Save parents[0].type only once
• Early returns in the decorator loop; avoids further iteration if not necessary.

This should speed up your function, especially under heavy use!

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

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


class FunctionKind(Enum):
    FUNCTION = "function"
    INSTANCE_METHOD = "instance_method"
    CLASS_METHOD = "class_method"
    STATIC_METHOD = "static_method"

class FunctionParent:
    def __init__(self, type_: str):
        self.type = type_
from codeflash.code_utils.static_analysis import function_kind


# Helper to parse function source code and return ast.FunctionDef or ast.AsyncFunctionDef
def get_func_node(src: str):
    module = ast.parse(src)
    for node in ast.walk(module):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            return node
    raise ValueError("No function node found in source.")

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

# 1. Basic Test Cases

def test_function_at_module_level():
    # Basic: function at module level (no parents)
    node = get_func_node("def foo(): pass")
    parents = []
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_function_with_function_parent():
    # Basic: function nested inside another function
    node = get_func_node("def bar():\n    def inner(): pass")
    # simulate being inside a function
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node.body[0], parents) # 583ns -> 666ns

def test_function_with_asyncfunction_parent():
    # Basic: function nested inside async function
    node = get_func_node("async def bar():\n    def inner(): pass")
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node.body[0], parents) # 583ns -> 583ns







def test_no_parents_none():
    # Edge: parents is None (should treat as no parents)
    node = get_func_node("def foo(): pass")
    codeflash_output = function_kind(node, None) # 542ns -> 833ns






def test_classdef_parent_not_first():
    # Edge: multiple parents, ClassDef not first
    node = get_func_node("def foo(): pass")
    parents = [FunctionParent("FunctionDef"), FunctionParent("ClassDef")]
    # Only first parent matters, so should be FUNCTION
    codeflash_output = function_kind(node, parents) # 708ns -> 1.25μs

def test_unknown_parent_type():
    # Edge: unknown parent type, should return None
    node = get_func_node("def foo(): pass")
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 500ns -> 500ns



def test_async_function_at_module_level():
    # Edge: async function at module level
    node = get_func_node("async def foo(): pass")
    parents = []
    codeflash_output = function_kind(node, parents) # 500ns -> 667ns

def test_function_with_weird_parent():
    # Edge: parent type is an unexpected string
    node = get_func_node("def foo(): pass")
    parents = [FunctionParent("WeirdParent")]
    codeflash_output = function_kind(node, parents) # 500ns -> 625ns

def test_function_with_non_list_parents():
    # Edge: parents is not a list (e.g., tuple)
    node = get_func_node("def foo(): pass")
    parents = tuple()
    codeflash_output = function_kind(node, parents) # 500ns -> 458ns


def test_many_methods_in_class():
    # Large scale: class with many methods, all instance methods
    src = "class C:\n" + "\n".join([f"    def foo{i}(self): pass" for i in range(500)])
    class_node = ast.parse(src).body[0]
    parents = [FunctionParent("ClassDef")]
    for func_node in class_node.body:
        codeflash_output = function_kind(func_node, parents) # 208ns -> 208ns


def test_large_nested_functions():
    # Large scale: deeply nested functions (up to 10 levels)
    src = ""
    indent = ""
    for i in range(10):
        src += f"{indent}def f{i}():\n"
        indent += "    "
    src += f"{indent}pass"
    node = ast.parse(src)
    # Go down to the innermost function
    func_node = node.body[0]
    parents = [FunctionParent("FunctionDef")] * 10
    # The innermost function node is at depth 10
    for _ in range(9):
        func_node = func_node.body[0]
    codeflash_output = function_kind(func_node, parents) # 750ns -> 875ns

def test_large_number_of_parents():
    # Large scale: parents list with many elements, first is ClassDef
    node = get_func_node("def foo(self): pass")
    parents = [FunctionParent("ClassDef")] + [FunctionParent("FunctionDef")] * 999
    codeflash_output = function_kind(node, parents) # 709ns -> 708ns

def test_large_number_of_module_level_functions():
    # Large scale: many module-level functions
    src = "\n".join([f"def foo{i}(): pass" for i in range(1000)])
    module = ast.parse(src)
    parents = []
    for func_node in module.body:
        codeflash_output = function_kind(func_node, parents) # 83ns -> 83ns
# 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 Any

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

# --- Mocking FunctionKind and FunctionParent for testability ---


class FunctionKind(Enum):
    FUNCTION = "function"
    CLASS_METHOD = "class_method"
    STATIC_METHOD = "static_method"
    INSTANCE_METHOD = "instance_method"

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 or AsyncFunctionDef node with decorators
def make_func_node(
    name: str = "f",
    decorators: list[Any] = None,
    async_: bool = False
) -> ast.FunctionDef | ast.AsyncFunctionDef:
    if decorators is None:
        decorators = []
    args = ast.arguments(
        posonlyargs=[],
        args=[],
        kwonlyargs=[],
        kw_defaults=[],
        defaults=[]
    )
    if async_:
        return ast.AsyncFunctionDef(
            name=name,
            args=args,
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )
    else:
        return ast.FunctionDef(
            name=name,
            args=args,
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )

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

def test_top_level_function_returns_functionkind_function():
    # Function at module level, no parents
    node = make_func_node()
    parents = []
    codeflash_output = function_kind(node, parents); result = codeflash_output # 417ns -> 458ns

def test_function_inside_function_returns_functionkind_function():
    # Function nested inside another function
    node = make_func_node()
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 625ns -> 709ns

def test_async_function_top_level():
    node = make_func_node(async_=True)
    parents = []
    codeflash_output = function_kind(node, parents); result = codeflash_output # 417ns -> 417ns

def test_class_instance_method():
    # Function inside a class, no decorators
    node = make_func_node()
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 708ns -> 750ns

def test_class_classmethod():
    # Function inside a class, with @classmethod
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 1.04μs

def test_class_staticmethod():
    # Function inside a class, with @staticmethod
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 834ns -> 875ns

def test_class_multiple_decorators_staticmethod_last():
    # Function inside a class, multiple decorators, staticmethod last
    decorator1 = ast.Name(id="something", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 1.04μs -> 1.08μs

def test_class_multiple_decorators_classmethod_first():
    # Function inside a class, multiple decorators, classmethod first
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="other", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 791ns -> 833ns

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

def test_empty_parents_returns_function():
    # No parents, should be FUNCTION
    node = make_func_node()
    codeflash_output = function_kind(node, []); result = codeflash_output # 417ns -> 459ns

def test_parent_is_unknown_type_returns_none():
    # Parent is not ClassDef or FunctionDef/AsyncFunctionDef
    node = make_func_node()
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 458ns -> 542ns

def test_decorator_is_not_ast_name():
    # Decorator is not ast.Name (e.g., ast.Call for @decorator())
    decorator = ast.Call(
        func=ast.Name(id="classmethod", ctx=ast.Load()),
        args=[],
        keywords=[]
    )
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 958ns

def test_decorator_is_ast_attribute():
    # Decorator is ast.Attribute (e.g., @abc.abstractmethod)
    decorator = ast.Attribute(
        value=ast.Name(id="abc", ctx=ast.Load()),
        attr="abstractmethod",
        ctx=ast.Load()
    )
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 750ns -> 792ns

def test_class_with_both_class_and_static_decorators():
    # Both classmethod and staticmethod present, classmethod first
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 792ns -> 833ns

def test_class_with_both_static_and_class_decorators():
    # Both staticmethod and classmethod, staticmethod first
    decorator1 = ast.Name(id="staticmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 834ns

def test_class_with_nonstring_decorator_id():
    # Decorator with id that is not a recognized string
    decorator = ast.Name(id="otherdecorator", ctx=ast.Load())
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 833ns -> 875ns

def test_async_function_inside_class():
    # Async function inside class, no decorators
    node = make_func_node(async_=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 667ns -> 666ns

def test_function_with_multiple_parents():
    # Multiple parents, only first is checked
    node = make_func_node()
    parents = [FunctionParent("ClassDef"), FunctionParent("Module")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 625ns -> 667ns

def test_function_with_parent_asyncfunctiondef():
    # Parent is AsyncFunctionDef
    node = make_func_node()
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 625ns -> 583ns

def test_function_with_parent_none_type():
    # Parent with type None
    class DummyParent:
        def __init__(self):
            self.type = None
    node = make_func_node()
    parents = [DummyParent()]
    codeflash_output = function_kind(node, parents); result = codeflash_output

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

def test_many_functions_in_module():
    # Test many top-level functions
    for i in range(100):
        node = make_func_node(name=f"func_{i}")
        parents = []
        codeflash_output = function_kind(node, parents); result = codeflash_output # 125ns -> 83ns

def test_many_class_methods():
    # Test many @classmethods in a class
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    for i in range(100):
        node = make_func_node(name=f"class_func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents); result = codeflash_output # 291ns -> 250ns

def test_many_static_methods():
    # Test many @staticmethods in a class
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    for i in range(100):
        node = make_func_node(name=f"static_func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents); result = codeflash_output # 291ns -> 250ns

def test_many_instance_methods_with_varied_decorators():
    # Test many instance methods with non-matching decorators
    for i in range(100):
        decorator = ast.Name(id=f"decorator_{i}", ctx=ast.Load())
        node = make_func_node(name=f"instance_func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents); result = codeflash_output # 291ns -> 250ns

def test_large_decorator_list_with_classmethod_at_end():
    # Function with 999 unrelated decorators and classmethod at the end
    decorators = [ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(999)]
    decorators.append(ast.Name(id="classmethod", ctx=ast.Load()))
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 59.8μs -> 52.8μs

def test_large_decorator_list_with_staticmethod_at_start():
    # Function with staticmethod as first of 1000 decorators
    decorators = [ast.Name(id="staticmethod", ctx=ast.Load())] + [
        ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(999)
    ]
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 875ns

def test_large_number_of_parents():
    # Only first parent should matter, even if there are many
    node = make_func_node()
    parents = [FunctionParent("ClassDef")] + [FunctionParent("Module") for _ in range(999)]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 709ns -> 708ns

def test_large_number_of_functions_varied_contexts():
    # Mix of top-level, nested, and class methods
    for i in range(50):
        node = make_func_node(name=f"func_{i}")
        codeflash_output = function_kind(node, []) # 125ns -> 125ns
        codeflash_output = function_kind(node, [FunctionParent("FunctionDef")]) # 125ns -> 125ns
        codeflash_output = function_kind(node, [FunctionParent("ClassDef")]) # 125ns -> 125ns
        decorator = ast.Name(id="classmethod", ctx=ast.Load())
        node2 = make_func_node(name=f"class_func_{i}", decorators=[decorator])
        codeflash_output = function_kind(node2, [FunctionParent("ClassDef")]) # 125ns -> 125ns
# 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-mbsmr7uy and push.