⚡️ Speed up function _dtype_to_init_repr_struct by 45%

[codeflash-ai]

📄 45% (0.45x) speedup for _dtype_to_init_repr_struct in py-polars/src/polars/datatypes/_utils.py

⏱️ Runtime : 830 microseconds → 574 microseconds (best of 241 runs)

📝 Explanation and details

The optimization achieves a 44% speedup through two key performance improvements:

1. Type-based dispatch optimization:

• Replaces isinstance(dtype, List) checks with direct type comparisons type(dtype) is List
• Direct type checks are faster because they avoid the inheritance chain traversal that isinstance() performs
• Early returns eliminate unnecessary variable assignments and reduce function overhead

2. Struct processing efficiency:

• Avoids the costly dict(dtype) conversion by calling dtype.items() directly when available
• Uses a pre-allocated list with append() instead of a list comprehension, which is more memory-efficient for complex nested structures
• Reduces temporary object creation during string formatting

Performance characteristics:

• Most effective for nested struct processing where recursive calls to dtype_to_init_repr are frequent
• Particularly beneficial for complex type hierarchies with multiple levels of List/Array/Struct nesting
• The type dispatch optimization helps most when processing mixed collections of different Polars data types

The test results show consistent performance gains across various scenarios, from simple single-field structs to deeply nested structures with 1000+ fields, indicating the optimizations scale well with complexity.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from polars.datatypes._utils import _dtype_to_init_repr_struct


# Minimal stand-in classes to simulate Polars DataTypes for testing
class DummyDataType:
    """A dummy base class for simulating PolarsDataType."""
    def __repr__(self):
        return self.__class__.__name__

class Int64(DummyDataType):
    pass

class Utf8(DummyDataType):
    pass

class Boolean(DummyDataType):
    pass

class Float32(DummyDataType):
    pass

class List(DummyDataType):
    def __init__(self, inner=None):
        self.inner = inner
    def __repr__(self):
        return f"List({self.inner!r})" if self.inner else "List()"

class Array(DummyDataType):
    def __init__(self, inner=None, shape=()):
        self.inner = inner
        self.shape = shape
    def __repr__(self):
        return f"Array({self.inner!r}, shape={self.shape})"

class Struct(DummyDataType):
    def __init__(self, fields=None):
        self._fields = fields or {}
    def __iter__(self):
        return iter(self._fields.items())
    def __getitem__(self, key):
        return self._fields[key]
    def __len__(self):
        return len(self._fields)
    def __repr__(self):
        return f"Struct({self._fields!r})"
    def items(self):
        return self._fields.items()
    def __eq__(self, other):
        return isinstance(other, Struct) and self._fields == other._fields

def _dtype_to_init_repr_list(dtype: List, prefix: str) -> str:
    class_name = dtype.__class__.__name__
    if dtype.inner is not None:
        inner_repr = dtype_to_init_repr(dtype.inner, prefix)
    else:
        inner_repr = ""
    init_repr = f"{prefix}{class_name}({inner_repr})"
    return init_repr

def _dtype_to_init_repr_array(dtype: Array, prefix: str) -> str:
    class_name = dtype.__class__.__name__
    if dtype.inner is not None:
        inner_repr = dtype_to_init_repr(dtype.inner, prefix)
    else:
        inner_repr = ""
    init_repr = f"{prefix}{class_name}({inner_repr}, shape={dtype.shape})"
    return init_repr
from polars.datatypes._utils import _dtype_to_init_repr_struct

# unit tests

# 1. Basic Test Cases

def test_struct_empty_fields():
    # Test Struct with no fields
    dtype = Struct({})
    expected = "pl.Struct({})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_single_field_primitive():
    # Test Struct with one primitive field
    dtype = Struct({'a': Int64()})
    expected = "pl.Struct({'a': pl.Int64})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_multiple_fields_primitives():
    # Test Struct with multiple primitive fields
    dtype = Struct({'a': Int64(), 'b': Utf8(), 'c': Boolean()})
    expected = "pl.Struct({'a': pl.Int64, 'b': pl.Utf8, 'c': pl.Boolean})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")



def test_struct_with_struct_field():
    # Test Struct with a nested Struct field
    dtype = Struct({
        'a': Struct({'x': Int64(), 'y': Utf8()}),
        'b': Boolean()
    })
    expected = "pl.Struct({'a': pl.Struct({'x': pl.Int64, 'y': pl.Utf8}), 'b': pl.Boolean})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_mixed_nested_fields():
    # Test Struct with nested List, Array, and Struct fields
    dtype = Struct({
        'a': List(Array(Int64(), shape=(2,))),
        'b': Struct({'z': Utf8()}),
        'c': Boolean()
    })
    expected = (
        "pl.Struct({'a': pl.List(pl.Array(pl.Int64, shape=(2,))), "
        "'b': pl.Struct({'z': pl.Utf8}), 'c': pl.Boolean})"
    )
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_custom_prefix():
    # Test Struct with a custom prefix
    dtype = Struct({'a': Int64(), 'b': Utf8()})
    expected = "my_prefix.Struct({'a': my_prefix.Int64, 'b': my_prefix.Utf8})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "my_prefix.")

# 2. Edge Test Cases

def test_struct_field_name_edge_cases():
    # Test Struct with unusual field names
    dtype = Struct({
        '': Int64(),         # empty string
        'with space': Utf8(),
        '123': Boolean(),    # numeric string
        '!@#': Float32(),    # special chars
    })
    expected = (
        "pl.Struct({'': pl.Int64, 'with space': pl.Utf8, '123': pl.Boolean, '!@#': pl.Float32})"
    )
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_none_inner_dtype():
    # Test Struct with a field whose dtype is None (should still represent as None)
    dtype = Struct({'a': None})
    expected = "pl.Struct({'a': pl.None})"
    # To handle None, we patch dtype_to_init_repr to handle None for this test
    def patched_dtype_to_init_repr(dtype, prefix="pl."):
        if dtype is None:
            return f"{prefix}None"
        return dtype_to_init_repr(dtype, prefix)
    # Patch only for this test
    class PatchedStruct(Struct):
        pass
    PatchedStruct.items = lambda self: self._fields.items()
    # We patch _dtype_to_init_repr_struct locally
    def patched_struct(dtype, prefix):
        class_name = dtype.__class__.__name__
        items = dtype.items()
        inner_list = [
            f"{field_name!r}: {patched_dtype_to_init_repr(inner_dtype, prefix)}"
            for field_name, inner_dtype in items
        ]
        inner_repr = "{" + ", ".join(inner_list) + "}"
        return f"{prefix}{class_name}({inner_repr})"

def test_struct_with_duplicate_field_names():
    # Test Struct with duplicate field names (should only keep one, as dict keys are unique)
    dtype = Struct({'a': Int64(), 'a': Utf8()})
    expected = "pl.Struct({'a': pl.Utf8})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_non_string_field_names():
    # Test Struct with non-string field names (e.g., int, tuple)
    dtype = Struct({1: Int64(), (2, 3): Utf8()})
    expected = "pl.Struct({1: pl.Int64, (2, 3): pl.Utf8})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_order_preservation():
    # Test that field order is preserved
    dtype = Struct({'z': Int64(), 'a': Utf8(), 'm': Boolean()})
    expected = "pl.Struct({'z': pl.Int64, 'a': pl.Utf8, 'm': pl.Boolean})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_empty_string_prefix():
    # Test Struct with empty string prefix
    dtype = Struct({'a': Int64()})
    expected = "Struct({'a': Int64})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "")


def test_struct_with_many_fields():
    # Test Struct with a large number of fields (up to 1000)
    num_fields = 1000
    fields = {f"f{i}": Int64() if i % 2 == 0 else Utf8() for i in range(num_fields)}
    dtype = Struct(fields)
    # Build expected string
    inner = ", ".join(
        [f"'f{i}': pl.Int64" if i % 2 == 0 else f"'f{i}': pl.Utf8" for i in range(num_fields)]
    )
    expected = f"pl.Struct({{{inner}}})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_deeply_nested_structs():
    # Test Struct with nested Structs up to depth 10
    dtype = Int64()
    for i in range(10):
        dtype = Struct({f"level_{i}": dtype})
    # Build expected string
    expected = "pl.Struct(" + "".join([f"{{'level_{i}': " for i in range(9, -1, -1)]) + "pl.Int64" + "}" * 10 + ")"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_large_nested_lists_and_arrays():
    # Test Struct with nested Lists and Arrays, each with 100 fields
    num_fields = 100
    fields = {
        f"l{i}": List(Array(Int64(), shape=(i+1,)))
        for i in range(num_fields)
    }
    dtype = Struct(fields)
    # Build expected string
    inner = ", ".join(
        [f"'l{i}': pl.List(pl.Array(pl.Int64, shape=({i+1},)))" for i in range(num_fields)]
    )
    expected = f"pl.Struct({{{inner}}})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")

def test_struct_with_various_types_large():
    # Test Struct with 500 fields of mixed types (Int64, Utf8, Boolean, List, Array)
    num_fields = 500
    types = [Int64, Utf8, Boolean, lambda: List(Int64()), lambda: Array(Utf8(), shape=(2,))]
    fields = {
        f"f{i}": types[i % len(types)]() if callable(types[i % len(types)]) else types[i % len(types)]()
        for i in range(num_fields)
    }
    dtype = Struct(fields)
    # Build expected string
    def type_str(i):
        t = types[i % len(types)]
        if t is Int64:
            return "pl.Int64"
        elif t is Utf8:
            return "pl.Utf8"
        elif t is Boolean:
            return "pl.Boolean"
        elif callable(t):
            if t.__name__ == "<lambda>":
                # Check which lambda: List(Int64) or Array(Utf8)
                if i % len(types) == 3:
                    return "pl.List(pl.Int64)"
                elif i % len(types) == 4:
                    return "pl.Array(pl.Utf8, shape=(2,))"
        return "pl.Unknown"
    inner = ", ".join([f"'f{i}': {type_str(i)}" for i in range(num_fields)])
    expected = f"pl.Struct({{{inner}}})"
    codeflash_output = _dtype_to_init_repr_struct(dtype, "pl.")
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest
from polars.datatypes._utils import _dtype_to_init_repr_struct


# Minimal mock classes to simulate PolarsDataType, List, Array, Struct
class PolarsDataType:
    pass

class Struct(PolarsDataType, dict):
    def __init__(self, fields=None):
        # fields: dict mapping field_name -> dtype
        super().__init__()
        if fields:
            for k, v in fields.items():
                self[k] = v
from polars.datatypes._utils import _dtype_to_init_repr_struct

# ---- UNIT TESTS ----

# Basic test: Struct with one field, primitive type
def test_struct_one_field_primitive():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    s = Struct({'a': Int32()})
    # Should produce: pl.Struct({'a': pl.Int32})
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Basic test: Struct with multiple primitive fields
def test_struct_multiple_fields_primitive():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    class Float64(PolarsDataType):
        def __repr__(self): return "Float64"
    s = Struct({'a': Int32(), 'b': Float64()})
    # Order in dict is preserved in Python 3.7+
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Basic test: Struct with nested List field

def test_struct_with_struct_field():
    class Utf8(PolarsDataType):
        def __repr__(self): return "Utf8"
    inner_struct = Struct({'x': Utf8()})
    s = Struct({'outer': inner_struct})
    # Should produce: pl.Struct({'outer': pl.Struct({'x': pl.Utf8})})
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Basic test: Struct with List of Structs

def test_struct_empty():
    s = Struct({})
    # Should produce: pl.Struct({})
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Edge case: Struct with field names that are not valid identifiers
def test_struct_with_strange_field_names():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    s = Struct({'1foo': Int32(), 'bar-baz': Int32()})
    # Should quote field names
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Edge case: Struct with None as a field dtype
def test_struct_with_none_dtype():
    s = Struct({'foo': None})
    # Should print None as pl.None
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Edge case: Struct with deeply nested structures
def test_struct_deeply_nested():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    deep = Struct({'a': Struct({'b': Struct({'c': Int32()})})})
    expected = "pl.Struct({'a': pl.Struct({'b': pl.Struct({'c': pl.Int32})})})"
    codeflash_output = _dtype_to_init_repr_struct(deep, "pl.")

# Edge case: Struct with Array fields

def test_struct_with_empty_field_name():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    s = Struct({'': Int32()})
    # Should quote empty string as field name
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Edge case: Struct with non-string field names (should still quote them)
def test_struct_with_nonstring_field_name():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    s = Struct({42: Int32()})

# Large scale: Struct with 100 fields
def test_struct_many_fields():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    fields = {f'col{i}': Int32() for i in range(100)}
    s = Struct(fields)
    codeflash_output = _dtype_to_init_repr_struct(s, "pl."); result = codeflash_output
    # Check that all fields are present and correctly formatted
    for i in range(100):
        pass

# Large scale: Struct with nested Structs (10 levels deep)
def test_struct_deep_nesting():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    s = Int32()
    for i in range(10):
        s = Struct({f'level{i}': s})
    # Should not raise or hang
    codeflash_output = _dtype_to_init_repr_struct(s, "pl."); result = codeflash_output
    # Check that all levels are present
    for i in range(10):
        pass

# Large scale: Struct with 100 fields, each a List of Structs

def test_struct_with_struct_with_empty_struct():
    empty_struct = Struct({})
    s = Struct({'empty': empty_struct})
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Edge case: Struct with a field whose dtype is itself (recursive, should not infinite loop)
def test_struct_recursive_reference():
    # This is a pathological case; let's see if it handles gracefully
    s = Struct({})
    s['self'] = s
    # Should not infinite loop; should handle recursion depth
    try:
        codeflash_output = _dtype_to_init_repr_struct(s, "pl."); result = codeflash_output
    except RecursionError:
        pytest.skip("RecursionError: recursive struct not supported (acceptable)")

# Edge case: Struct with non-PolarsDataType field (should still call repr)
def test_struct_with_non_polars_dtype():
    class Dummy:
        def __repr__(self): return "Dummy"
    s = Struct({'foo': Dummy()})
    codeflash_output = _dtype_to_init_repr_struct(s, "pl.")

# Edge case: Custom prefix
def test_struct_with_custom_prefix():
    class Int32(PolarsDataType):
        def __repr__(self): return "Int32"
    s = Struct({'a': Int32()})
    codeflash_output = _dtype_to_init_repr_struct(s, "mypl.")
# 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-_dtype_to_init_repr_struct-mgyu9mnu and push.