⚡️ Speed up method AES._pad by 23%

[codeflash-ai]

📄 23% (0.23x) speedup for AES._pad in skyvern/forge/sdk/encrypt/aes.py

⏱️ Runtime : 89.2 microseconds → 72.8 microseconds (best of 250 runs)

📝 Explanation and details

The optimization replaces the inefficient padding creation pattern with a more direct approach, achieving a 22% speedup by eliminating unnecessary list allocation.

Key Change:

• Original: padding = bytes([padding_length] * padding_length) followed by return data + padding
• Optimized: return data + bytes([padding_length]) * padding_length

Why This Is Faster:
The original code creates an intermediate list [padding_length] * padding_length (e.g., [5, 5, 5, 5, 5] for padding_length=5) and then converts it to bytes. This involves:

1. List allocation and population
2. Conversion from list to bytes object
3. Separate concatenation operation

The optimized version uses bytes([padding_length]) * padding_length, which:

1. Creates a single-byte bytes object bytes([padding_length])
2. Uses Python's efficient bytes multiplication to repeat it
3. Performs concatenation in the same expression

This eliminates the intermediate list allocation and leverages Python's optimized bytes repetition, which is implemented in C and avoids Python-level iteration.

Performance Impact:
The line profiler shows the padding creation line improved significantly - the total function time decreased from 255.4μs to 213.2μs. Test results demonstrate consistent 10-35% speedups across various input sizes, with particularly strong improvements for edge cases like empty inputs (34.4% faster) and small data sizes.

Test Case Benefits:
The optimization performs well across all scenarios:

• Small inputs (1-16 bytes): 12-35% faster
• Large inputs (999+ bytes): 12-22% faster
• Edge cases (empty, block-aligned): 22-35% faster

This suggests the optimization is universally beneficial for the PKCS#7 padding operation regardless of input characteristics.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import hashlib

# imports
import pytest  # used for our unit tests
from skyvern.forge.sdk.encrypt.aes import AES

default_iv = hashlib.md5(b"deterministic_iv_0123456789").digest()
default_salt = hashlib.md5(b"deterministic_salt_0123456789").digest()
from skyvern.forge.sdk.encrypt.aes import AES

# unit tests

@pytest.fixture
def aes():
    # Provide a default AES instance for tests
    return AES(secret_key="testkey")

# 1. Basic Test Cases

def test_pad_empty_bytes(aes):
    # Test padding for empty input
    codeflash_output = aes._pad(b""); padded = codeflash_output # 1.39μs -> 1.03μs (34.4% faster)

def test_pad_exact_block_size(aes):
    # Test padding for data exactly one block long
    data = b"A" * 16
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.17μs -> 984ns (19.3% faster)

def test_pad_one_byte_short(aes):
    # Test padding for data one byte short of block size
    data = b"B" * 15
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.03μs -> 920ns (12.0% faster)

def test_pad_one_byte_over_block(aes):
    # Test padding for data one byte over block size
    data = b"C" * 17
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.17μs -> 890ns (30.9% faster)

def test_pad_arbitrary_length(aes):
    # Test padding for arbitrary length data
    data = b"hello world"
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.08μs -> 881ns (22.1% faster)
    pad_len = 16 - (len(data) % 16)

# 2. Edge Test Cases

def test_pad_max_block_minus_one(aes):
    # Test data length 15 (just under block size)
    data = b"D" * 15
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.02μs -> 909ns (12.5% faster)

def test_pad_multiple_blocks(aes):
    # Test data length is multiple of block size
    data = b"E" * 32
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.11μs -> 896ns (23.8% faster)

def test_pad_block_size_minus_one_multiple(aes):
    # Test data length is one less than multiple of block size (e.g., 31)
    data = b"F" * 31
    codeflash_output = aes._pad(data); padded = codeflash_output # 950ns -> 860ns (10.5% faster)

def test_pad_all_byte_values(aes):
    # Test padding for data containing all byte values
    data = bytes(range(256))
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.11μs -> 910ns (22.3% faster)
    pad_len = 16 - (len(data) % 16)

def test_pad_zero_length_block_multiple(aes):
    # Test padding for zero-length data (edge case, already covered above)
    codeflash_output = aes._pad(b""); padded = codeflash_output # 1.18μs -> 963ns (22.5% faster)

def test_pad_large_block_minus_one(aes):
    # Test padding for 255 bytes (not a multiple of block size)
    data = b"G" * 255
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.00μs -> 872ns (14.7% faster)
    pad_len = 16 - (len(data) % 16)

def test_pad_with_non_ascii_bytes(aes):
    # Test padding for data with non-ASCII bytes
    data = b"\xff\xfe\xfd\xfc" * 4  # 16 bytes
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.09μs -> 863ns (26.1% faster)

# 3. Large Scale Test Cases

def test_pad_large_data_exact_block(aes):
    # Test large data that is an exact multiple of block size
    data = b"H" * 16 * 50  # 800 bytes
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.22μs -> 1.07μs (14.5% faster)

def test_pad_large_data_non_block_multiple(aes):
    # Test large data that is not a multiple of block size
    data = b"I" * 999
    pad_len = 16 - (len(data) % 16)
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.26μs -> 1.12μs (12.5% faster)

def test_pad_large_data_block_minus_one(aes):
    # Test large data that is one less than a block multiple
    data = b"J" * (16 * 62 - 1)  # 991 bytes
    pad_len = 1
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.12μs -> 952ns (17.2% faster)

def test_pad_max_allowed_size(aes):
    # Test padding for maximum allowed size (999 bytes)
    data = b"K" * 999
    pad_len = 16 - (len(data) % 16)
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.24μs -> 1.06μs (16.3% faster)

def test_pad_repeated_patterns(aes):
    # Test padding for data with repeated patterns
    data = (b"0123456789ABCDEF" * 62)[:999]  # 999 bytes
    pad_len = 16 - (len(data) % 16)
    codeflash_output = aes._pad(data); padded = codeflash_output # 1.30μs -> 1.09μs (19.5% faster)

# Additional: Test that padding bytes are correct for all possible padding lengths
@pytest.mark.parametrize("data_len", range(0, 32))
def test_pad_padding_byte_value(aes, data_len):
    # Test that the padding byte value is correct for all possible lengths
    data = b"L" * data_len
    codeflash_output = aes._pad(data); padded = codeflash_output # 33.5μs -> 26.8μs (24.9% faster)
    pad_len = 16 - (data_len % 16)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import hashlib

# imports
import pytest  # used for our unit tests
from skyvern.forge.sdk.encrypt.aes import AES

default_iv = hashlib.md5(b"deterministic_iv_0123456789").digest()
default_salt = hashlib.md5(b"deterministic_salt_0123456789").digest()
from skyvern.forge.sdk.encrypt.aes import AES

# unit tests

# ===== Basic Test Cases =====

def test_pad_empty_bytes():
    """Test padding for empty input (should pad to one full block)."""
    aes = AES(secret_key="test")
    codeflash_output = aes._pad(b""); result = codeflash_output # 1.31μs -> 1.04μs (26.8% faster)

def test_pad_exact_block_size():
    """Test padding for input with exact block size (should add a full block of padding)."""
    aes = AES(secret_key="test")
    data = b"A" * 16
    codeflash_output = aes._pad(data); result = codeflash_output # 1.10μs -> 894ns (23.5% faster)

def test_pad_one_byte_short():
    """Test padding for input one byte short of block size."""
    aes = AES(secret_key="test")
    data = b"A" * 15
    codeflash_output = aes._pad(data); result = codeflash_output # 995ns -> 872ns (14.1% faster)

def test_pad_various_lengths():
    """Test padding for various lengths less than block size."""
    aes = AES(secret_key="test")
    for length in range(1, 16):
        data = b"B" * length
        codeflash_output = aes._pad(data); result = codeflash_output # 5.64μs -> 4.22μs (33.7% faster)
        pad_len = 16 - length

def test_pad_multiple_blocks():
    """Test padding for input with multiple blocks and extra bytes."""
    aes = AES(secret_key="test")
    data = b"C" * 30  # 1 full block + 14 bytes
    codeflash_output = aes._pad(data); result = codeflash_output # 896ns -> 759ns (18.1% faster)
    pad_len = 2

# ===== Edge Test Cases =====

def test_pad_block_size_minus_one():
    """Test padding for input of block_size - 1 (15 bytes)."""
    aes = AES(secret_key="test")
    data = b"D" * 15
    codeflash_output = aes._pad(data); result = codeflash_output # 874ns -> 815ns (7.24% faster)

def test_pad_maximum_single_block():
    """Test padding for input of 15 bytes (maximum before new block)."""
    aes = AES(secret_key="test")
    data = b"E" * 15
    codeflash_output = aes._pad(data); result = codeflash_output # 868ns -> 780ns (11.3% faster)

def test_pad_block_size_plus_one():
    """Test padding for input of block_size + 1 (17 bytes)."""
    aes = AES(secret_key="test")
    data = b"F" * 17
    codeflash_output = aes._pad(data); result = codeflash_output # 999ns -> 762ns (31.1% faster)
    pad_len = 15

def test_pad_all_possible_padding_lengths():
    """Test all possible padding lengths from 1 to 16."""
    aes = AES(secret_key="test")
    for pad_len in range(1, 17):
        data = b"G" * (16 - pad_len)
        codeflash_output = aes._pad(data); result = codeflash_output # 6.10μs -> 4.46μs (36.8% faster)

def test_pad_non_ascii_bytes():
    """Test padding for non-ASCII input bytes."""
    aes = AES(secret_key="test")
    data = bytes([0xff, 0x00, 0x7f, 0x80])
    codeflash_output = aes._pad(data); result = codeflash_output # 852ns -> 685ns (24.4% faster)
    pad_len = 12

def test_pad_zero_length_input():
    """Test that zero-length input returns a full block of padding."""
    aes = AES(secret_key="test")
    codeflash_output = aes._pad(b""); result = codeflash_output # 1.04μs -> 852ns (21.6% faster)

def test_pad_large_padding_value():
    """Test padding for input where padding value is maximum (16)."""
    aes = AES(secret_key="test")
    data = b""
    codeflash_output = aes._pad(data); result = codeflash_output # 1.01μs -> 790ns (28.0% faster)

def test_pad_with_all_byte_values():
    """Test padding for input containing all byte values 0-255."""
    aes = AES(secret_key="test")
    data = bytes(range(256))
    codeflash_output = aes._pad(data); result = codeflash_output # 995ns -> 773ns (28.7% faster)
    pad_len = 16 - (256 % 16)

# ===== Large Scale Test Cases =====

def test_pad_large_input_exact_block():
    """Test padding for large input that is an exact multiple of block size."""
    aes = AES(secret_key="test")
    data = b"H" * 512  # 32 blocks
    codeflash_output = aes._pad(data); result = codeflash_output # 1.07μs -> 919ns (16.9% faster)

def test_pad_large_input_non_multiple_block():
    """Test padding for large input that is not a multiple of block size."""
    aes = AES(secret_key="test")
    data = b"I" * 999  # Not a multiple of 16
    pad_len = 16 - (999 % 16)
    codeflash_output = aes._pad(data); result = codeflash_output # 1.24μs -> 1.10μs (12.3% faster)

def test_pad_input_just_under_1000_bytes():
    """Test padding for input just under 1000 bytes."""
    aes = AES(secret_key="test")
    data = b"J" * 999
    pad_len = 16 - (999 % 16)
    codeflash_output = aes._pad(data); result = codeflash_output # 1.20μs -> 1.06μs (13.0% faster)

def test_pad_input_just_over_1000_bytes():
    """Test padding for input just over 1000 bytes."""
    aes = AES(secret_key="test")
    data = b"K" * 1001
    pad_len = 16 - (1001 % 16)
    codeflash_output = aes._pad(data); result = codeflash_output # 1.23μs -> 1.05μs (17.1% faster)

def test_pad_input_all_zeros_large():
    """Test padding for large input of all zeros."""
    aes = AES(secret_key="test")
    data = bytes([0] * 960)  # 60 blocks
    codeflash_output = aes._pad(data); result = codeflash_output # 1.02μs -> 922ns (10.4% faster)

def test_pad_performance_large_input():
    """Test that padding a large input is efficient and correct."""
    aes = AES(secret_key="test")
    data = b"L" * 1000
    pad_len = 16 - (1000 % 16)
    codeflash_output = aes._pad(data); result = codeflash_output # 1.23μs -> 998ns (22.9% faster)
    # Check that the function runs quickly (no explicit timing, but should not timeout)

# ===== Negative Tests =====

def test_pad_non_bytes_input_raises():
    """Test that non-bytes input raises TypeError."""
    aes = AES(secret_key="test")
    with pytest.raises(TypeError):
        aes._pad("not bytes") # 1.83μs -> 1.67μs (9.77% faster)

def test_pad_none_input_raises():
    """Test that None input raises TypeError."""
    aes = AES(secret_key="test")
    with pytest.raises(TypeError):
        aes._pad(None) # 1.26μs -> 1.20μs (4.49% faster)

# ===== Consistency Tests =====

def test_pad_is_deterministic():
    """Test that padding the same input always yields the same output."""
    aes = AES(secret_key="test")
    data = b"repeatable"
    codeflash_output = aes._pad(data); result1 = codeflash_output # 1.54μs -> 1.48μs (3.91% faster)
    codeflash_output = aes._pad(data); result2 = codeflash_output # 370ns -> 384ns (3.65% slower)

def test_pad_different_inputs_different_outputs():
    """Test that different inputs yield different padded outputs."""
    aes = AES(secret_key="test")
    data1 = b"abc"
    data2 = b"abcd"
    codeflash_output = aes._pad(data1); result1 = codeflash_output # 1.10μs -> 872ns (26.0% faster)
    codeflash_output = aes._pad(data2); result2 = codeflash_output # 475ns -> 318ns (49.4% faster)
# 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-AES._pad-mjab2rc4 and push.