⚡️ Speed up function function_kind by 55%

[codeflash-ai]

📄 55% (0.55x) speedup for function_kind in codeflash/code_utils/static_analysis.py

⏱️ Runtime : 1.11 milliseconds → 720 microseconds (best of 219 runs)

📝 Explanation and details

Here’s a faster, equivalent version. I eliminated the useless loop, streamlined the type checks to use set lookup, and combined conditionals where possible. Your core slowness came from the unnecessary loop and repeated attribute access.
All existing comments are preserved, but since the program had none, this is the cleaned, fast version:

Key runtime improvements:

• Removed the always-empty for-loop.
• Used a set for "FunctionDef", "AsyncFunctionDef" to speed up type check.
• Pulled .type access once and reuse.
• Short-circuited if-else for clarity and speed.

This will be notably faster especially under CPython, and the logic and structure remain identical.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 3051 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 dataclasses import dataclass
from enum import Enum

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


# Minimal stub for FunctionParent as per the signature
@dataclass
class FunctionParent:
    type: str

# Enum for FunctionKind as per function_kind's return values
class FunctionKind(Enum):
    FUNCTION = "function"
    CLASS_METHOD = "class_method"
    STATIC_METHOD = "static_method"
    INSTANCE_METHOD = "instance_method"
from codeflash.code_utils.static_analysis import function_kind


# Helper to create a function node with decorators
def make_func(name="f", 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
        )

# ----------- UNIT TESTS ------------

# Basic Test Cases

def test_top_level_function():
    # Function at module level (no parents)
    node = make_func()
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 333ns

def test_top_level_async_function():
    # Async function at module level
    node = make_func(async_=True)
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 333ns

def test_nested_function():
    # Function inside another function
    node = make_func()
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

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

def test_instance_method():
    # Function inside a class, no decorators
    node = make_func()
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.08μs -> 625ns

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

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

# Edge Test Cases

def test_function_with_multiple_decorators_static_first():
    # Function with multiple decorators, @staticmethod is first
    node = make_func(decorators=[ast.Name(id="staticmethod"), ast.Name(id="other")])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 792ns

def test_function_with_multiple_decorators_classmethod_last():
    # Function with multiple decorators, @classmethod is last
    node = make_func(decorators=[ast.Name(id="other"), ast.Name(id="classmethod")])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.42μs -> 917ns

def test_function_with_non_name_decorator():
    # Function with a decorator that is not ast.Name
    node = make_func(decorators=[ast.Attribute(value=ast.Name(id="foo"), attr="bar")])
    parents = [FunctionParent("ClassDef")]
    # Should be instance method since no recognized decorator
    codeflash_output = function_kind(node, parents) # 1.21μs -> 750ns

def test_function_with_unrelated_decorator():
    # Function with unrelated decorator
    node = make_func(decorators=[ast.Name(id="other")])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 792ns

def test_function_with_both_class_and_static_decorators():
    # Function with both @classmethod and @staticmethod; first match wins
    node = make_func(decorators=[ast.Name(id="classmethod"), ast.Name(id="staticmethod")])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 750ns

def test_function_with_both_static_and_class_decorators():
    # Function with both @staticmethod and @classmethod; first match wins
    node = make_func(decorators=[ast.Name(id="staticmethod"), ast.Name(id="classmethod")])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.21μs -> 750ns

def test_function_with_empty_parents():
    # Function with empty parents list
    node = make_func()
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 333ns

def test_function_with_unknown_parent_type():
    # Function with an unknown parent type
    node = make_func()
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 791ns -> 375ns

def test_function_with_long_parents_chain_class_first():
    # Function with a long parent chain, class at the top
    node = make_func()
    parents = [FunctionParent("ClassDef"), FunctionParent("FunctionDef"), FunctionParent("If")]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 583ns

def test_function_with_long_parents_chain_function_first():
    # Function with a long parent chain, function at the top
    node = make_func()
    parents = [FunctionParent("FunctionDef"), FunctionParent("ClassDef"), FunctionParent("If")]
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_function_with_no_decorators_in_class():
    # Function in a class with empty decorator list
    node = make_func(decorators=[])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.08μs -> 584ns

def test_function_with_parent_class_but_empty_decorator_list():
    # Function in a class with no decorators
    node = make_func()
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.04μs -> 625ns


def test_many_parents_chain_with_classdef():
    # Long chain of parents, class at the top
    node = make_func()
    parents = [FunctionParent("ClassDef")] + [FunctionParent("If") for _ in range(999)]
    codeflash_output = function_kind(node, parents) # 10.7μs -> 834ns

def test_many_parents_chain_with_functiondef():
    # Long chain of parents, function at the top
    node = make_func()
    parents = [FunctionParent("FunctionDef")] + [FunctionParent("If") for _ in range(999)]
    codeflash_output = function_kind(node, parents) # 542ns -> 583ns

def test_large_number_of_decorators_with_classmethod():
    # Many decorators, only one is classmethod
    decorators = [ast.Name(id=f"decorator{i}") for i in range(500)] + [ast.Name(id="classmethod")]
    node = make_func(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 30.5μs -> 25.9μs

def test_large_number_of_decorators_with_staticmethod():
    # Many decorators, only one is staticmethod
    decorators = [ast.Name(id=f"decorator{i}") for i in range(500)] + [ast.Name(id="staticmethod")]
    node = make_func(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 30.4μs -> 26.4μs

def test_large_number_of_decorators_none_matching():
    # Many unrelated decorators, should be instance method
    decorators = [ast.Name(id=f"decorator{i}") for i in range(1000)]
    node = make_func(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 57.5μs -> 50.3μs

def test_large_number_of_functions_module_level():
    # Many top-level functions, all should be FUNCTION
    for i in range(1000):
        node = make_func(name=f"f{i}")
        parents = []
        codeflash_output = function_kind(node, parents) # 83ns -> 83ns

def test_large_number_of_functions_in_class():
    # Many functions in a class, all should be INSTANCE_METHOD
    parents = [FunctionParent("ClassDef")]
    for i in range(1000):
        node = make_func(name=f"f{i}")
        codeflash_output = function_kind(node, parents) # 333ns -> 166ns
# 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, Optional

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


# Mocks for FunctionKind and FunctionParent (since codeflash.models.models is unavailable)
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


# Helper to build a function node with optional decorators
def make_func(name="f", async_=False, decorators=None):
    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
        )

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

def test_top_level_function():
    # Function at module level (no parents)
    node = make_func()
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 333ns

def test_top_level_async_function():
    # Async function at module level
    node = make_func(async_=True)
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 333ns

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

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

def test_instance_method():
    # Regular method in a class (no decorators)
    node = make_func()
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.08μs -> 625ns

def test_class_method():
    # Method in a class decorated with @classmethod
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 791ns

def test_static_method():
    # Method in a class decorated with @staticmethod
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 833ns

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

def test_function_with_multiple_parents():
    # Deeply nested: function inside a class inside a function
    node = make_func()
    parents = [FunctionParent("ClassDef"), FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents) # 1.08μs -> 625ns

def test_function_with_multiple_class_parents():
    # Function inside class inside class (e.g., inner class method)
    node = make_func()
    parents = [FunctionParent("ClassDef"), FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.08μs -> 583ns

def test_class_method_with_multiple_decorators():
    # @classmethod with other decorators
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="other_decorator", ctx=ast.Load())
    node = make_func(decorators=[decorator2, decorator1])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.46μs -> 958ns

def test_static_method_with_multiple_decorators():
    # @staticmethod with other decorators
    decorator1 = ast.Name(id="staticmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="other_decorator", ctx=ast.Load())
    node = make_func(decorators=[decorator2, decorator1])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.46μs -> 1.00μs

def test_method_with_non_name_decorator():
    # Decorator is not an ast.Name (e.g., ast.Attribute)
    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 decorator is not recognized
    codeflash_output = function_kind(node, parents) # 1.25μs -> 792ns

def test_method_with_unknown_decorator_name():
    # Decorator is an ast.Name but not classmethod/staticmethod
    decorator = ast.Name(id="custom", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 833ns

def test_function_with_unknown_parent_type():
    # Parent is not a function or class
    node = make_func()
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 791ns -> 375ns

def test_function_with_empty_parents_list():
    # Explicitly empty parent list
    node = make_func()
    parents = []
    codeflash_output = function_kind(node, parents) # 333ns -> 375ns

def test_function_with_none_parents():
    # Parents is None (should behave like empty list)
    node = make_func()
    parents = None
    codeflash_output = function_kind(node, parents or []) # 375ns -> 375ns

def test_function_with_long_parents_chain():
    # Long chain of parents, top is ClassDef
    node = make_func()
    parents = [FunctionParent("ClassDef")] + [FunctionParent("FunctionDef") for _ in range(10)]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 625ns

def test_function_with_long_non_class_parents_chain():
    # Long chain, but no ClassDef at top
    node = make_func()
    parents = [FunctionParent("Module")] + [FunctionParent("FunctionDef") for _ in range(10)]
    codeflash_output = function_kind(node, parents) # 875ns -> 416ns

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

def test_many_decorators_performance():
    # Function with many (non-matching) decorators, only last is classmethod
    decorators = [ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(998)]
    decorators.append(ast.Name(id="classmethod", ctx=ast.Load()))
    node = make_func(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 55.4μs -> 48.2μs

def test_many_functions_in_large_class():
    # Simulate 999 methods in a class, all should be instance methods (no decorators)
    parents = [FunctionParent("ClassDef")]
    for i in range(999):
        node = make_func(name=f"f{i}")
        codeflash_output = function_kind(node, parents) # 333ns -> 166ns

def test_many_nested_functions():
    # Deeply nested function (999 levels), top parent is FunctionDef
    node = make_func()
    parents = [FunctionParent("FunctionDef") for _ in range(999)]
    codeflash_output = function_kind(node, parents) # 459ns -> 583ns

def test_large_parents_chain_with_classdef_at_top():
    # 999 parents, top is ClassDef, rest are FunctionDef
    node = make_func()
    parents = [FunctionParent("ClassDef")] + [FunctionParent("FunctionDef") for _ in range(998)]
    codeflash_output = function_kind(node, parents) # 10.5μs -> 625ns

def test_large_parents_chain_with_non_classdef_at_top():
    # 999 parents, top is Module, rest are ClassDef
    node = make_func()
    parents = [FunctionParent("Module")] + [FunctionParent("ClassDef") for _ in range(998)]
    codeflash_output = function_kind(node, parents) # 10.2μs -> 375ns

# -----------------------
# Mutation-sensitive tests
# -----------------------

def test_staticmethod_vs_classmethod_priority():
    # If both @staticmethod and @classmethod are present, @classmethod wins if first
    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) # 583ns -> 417ns

    # If @staticmethod is first, it should still find @classmethod if present
    node2 = make_func(decorators=[decorator2, decorator1])
    codeflash_output = function_kind(node2, parents) # 583ns -> 417ns

def test_case_sensitivity_of_decorators():
    # Decorator with wrong case should not be recognized
    decorator = ast.Name(id="ClassMethod", ctx=ast.Load())
    node = make_func(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 875ns

def test_async_function_in_class():
    # Async method in class should be classified correctly (instance method if no decorator)
    node = make_func(async_=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.00μs -> 542ns

def test_async_classmethod():
    # Async method in class with @classmethod
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func(async_=True, decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 791ns

def test_async_staticmethod():
    # Async method in class with @staticmethod
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func(async_=True, decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 875ns
# 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-mbzomb0i and push.