⚡️ Speed up function _get_grid_subplot by 20%

[codeflash-ai]

📄 20% (0.20x) speedup for _get_grid_subplot in plotly/_subplots.py

⏱️ Runtime : 102 microseconds → 85.2 microseconds (best of 354 runs)

📝 Explanation and details

Here is an optimized version of your program.
Key optimizations.

• Avoid repeated attribute lookups (especially fig._grid_ref and fig.layout), store them in local variables.
• Calculate rows and cols only once.
• Avoid redundant ternaries and unnecessary branching.
• Use f-strings for marginally faster formatting.
• Inline trivial variables where safe and short.

The logic and behavior (including exception text and order) are unmodified.

Notes:

• All lookups for fig._grid_ref, fig.layout are minimized and reused.
• No function signatures or return values are changed.
• All comments are preserved (or untouched for unmodified code regions).
• Helper constructs and error messages are identical.
• Readability and maintainability are improved with minor formatting tweaks.
• No new dependencies are introduced.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import collections

# imports
import pytest  # used for our unit tests
from plotly._subplots import _get_grid_subplot

# function to test
# Constants
# ---------
# Subplot types that are each individually positioned with a domain
#
# Each of these subplot types has a `domain` property with `x`/`y`
# properties.
# Note that this set does not contain `xaxis`/`yaxis` because these behave a
# little differently.

# Named tuple to hold an xaxis/yaxis pair that represent a single subplot
SubplotXY = collections.namedtuple("SubplotXY", ("xaxis", "yaxis"))
SubplotDomain = collections.namedtuple("SubplotDomain", ("x", "y"))
from plotly._subplots import _get_grid_subplot

# --- Helper classes for mocking subplot structure ---

class MockSubplotRef:
    """Mocks a subplot reference object with layout_keys and trace_kwargs."""
    def __init__(self, layout_keys=None, trace_kwargs=None):
        self.layout_keys = layout_keys or []
        self.trace_kwargs = trace_kwargs or {}

class MockLayout(dict):
    """Mocks a layout object that supports dict access."""
    def __getitem__(self, key):
        return dict.__getitem__(self, key)

class MockFigure:
    """Mocks a figure object with _grid_ref and layout."""
    def __init__(self, grid_ref, layout):
        self._grid_ref = grid_ref
        self.layout = layout

# --- Unit tests ---

# BASIC TEST CASES

def test_single_subplot_with_one_layout_key():
    # 1x1 grid, layout_keys = ['xaxis1']
    layout = MockLayout({'xaxis1': 'xaxis-object'})
    grid_ref = [
        [[MockSubplotRef(layout_keys=['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    # Should return the layout object for 'xaxis1'
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output

def test_single_subplot_with_two_layout_keys():
    # 1x1 grid, layout_keys = ['xaxis1', 'yaxis1']
    layout = MockLayout({'xaxis1': 'xaxis-object', 'yaxis1': 'yaxis-object'})
    grid_ref = [
        [[MockSubplotRef(layout_keys=['xaxis1', 'yaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    # Should return SubplotXY with xaxis and yaxis
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output

def test_single_subplot_with_domain():
    # 1x1 grid, layout_keys = [], trace_kwargs has domain
    domain = {'x': [0, 0.5], 'y': [0.5, 1]}
    grid_ref = [
        [[MockSubplotRef(layout_keys=[], trace_kwargs={'domain': domain})]]
    ]
    layout = MockLayout({})
    fig = MockFigure(grid_ref, layout)
    # Should return SubplotDomain with x and y from domain
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output

def test_multiple_rows_and_columns_basic_access():
    # 2x2 grid, each with different layout_keys
    layout = MockLayout({
        'xaxis1': 'x1', 'yaxis1': 'y1',
        'xaxis2': 'x2', 'yaxis2': 'y2',
        'xaxis3': 'x3', 'yaxis3': 'y3',
        'xaxis4': 'x4', 'yaxis4': 'y4'
    })
    grid_ref = [
        [[MockSubplotRef(['xaxis1', 'yaxis1'])], [MockSubplotRef(['xaxis2', 'yaxis2'])]],
        [[MockSubplotRef(['xaxis3', 'yaxis3'])], [MockSubplotRef(['xaxis4', 'yaxis4'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    # Access each subplot and check correct mapping
    codeflash_output = _get_grid_subplot(fig, 1, 1)
    codeflash_output = _get_grid_subplot(fig, 1, 2)
    codeflash_output = _get_grid_subplot(fig, 2, 1)
    codeflash_output = _get_grid_subplot(fig, 2, 2)

# EDGE TEST CASES

def test_missing_grid_ref_attribute():
    # fig has no _grid_ref attribute
    class NoGridRef:
        pass
    fig = NoGridRef()
    with pytest.raises(Exception) as e:
        _get_grid_subplot(fig, 1, 1)

def test_row_argument_out_of_bounds_low():
    # row < 1
    layout = MockLayout({'xaxis1': 'x1'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 0, 1)

def test_row_argument_out_of_bounds_high():
    # row > number of rows
    layout = MockLayout({'xaxis1': 'x1'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 2, 1)

def test_col_argument_out_of_bounds_low():
    # col < 1
    layout = MockLayout({'xaxis1': 'x1'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 1, 0)

def test_col_argument_out_of_bounds_high():
    # col > number of columns
    layout = MockLayout({'xaxis1': 'x1'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 1, 2)

def test_row_argument_not_integer():
    # row is not an integer
    layout = MockLayout({'xaxis1': 'x1'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 'a', 1)

def test_col_argument_not_integer():
    # col is not an integer
    layout = MockLayout({'xaxis1': 'x1'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 1, None)

def test_empty_subplot_ref_returns_none():
    # grid_ref cell is an empty list
    layout = MockLayout({})
    grid_ref = [
        [[]]
    ]
    fig = MockFigure(grid_ref, layout)
    codeflash_output = _get_grid_subplot(fig, 1, 1)

def test_secondary_y_missing_returns_none():
    # secondary_y=True, but only one subplot_ref in cell
    layout = MockLayout({})
    grid_ref = [
        [[MockSubplotRef(['xaxis1'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    codeflash_output = _get_grid_subplot(fig, 1, 1, secondary_y=True)

def test_secondary_y_with_layout_keys():
    # secondary_y=True, second subplot_ref present
    layout = MockLayout({'yaxis2': 'y2'})
    grid_ref = [
        [[MockSubplotRef(['xaxis1']), MockSubplotRef(['yaxis2'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    codeflash_output = _get_grid_subplot(fig, 1, 1, secondary_y=True); result = codeflash_output

def test_unexpected_layout_keys_length():
    # layout_keys has more than 2 elements
    layout = MockLayout({'a': 1, 'b': 2, 'c': 3})
    grid_ref = [
        [[MockSubplotRef(['a', 'b', 'c'])]]
    ]
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(ValueError) as e:
        _get_grid_subplot(fig, 1, 1)

def test_domain_key_missing_in_trace_kwargs():
    # layout_keys empty, but trace_kwargs missing 'domain'
    grid_ref = [
        [[MockSubplotRef(layout_keys=[], trace_kwargs={})]]
    ]
    layout = MockLayout({})
    fig = MockFigure(grid_ref, layout)
    with pytest.raises(KeyError):
        _get_grid_subplot(fig, 1, 1)

# LARGE SCALE TEST CASES

def test_large_grid_access_first_and_last():
    # Test a 50x50 grid (2500 subplots), access first and last
    n = 50
    layout = MockLayout({f'xaxis{i*n+j+1}': f'x{i*n+j+1}' for i in range(n) for j in range(n)})
    grid_ref = []
    count = 1
    for i in range(n):
        row = []
        for j in range(n):
            row.append([MockSubplotRef([f'xaxis{count}'])])
            count += 1
        grid_ref.append(row)
    fig = MockFigure(grid_ref, layout)
    # First subplot
    codeflash_output = _get_grid_subplot(fig, 1, 1)
    # Last subplot
    codeflash_output = _get_grid_subplot(fig, n, n)

def test_large_grid_secondary_y():
    # 20x20 grid, each cell has two subplot_refs for secondary_y
    n = 20
    layout = MockLayout({f'xaxis{i*n+j+1}': f'x{i*n+j+1}' for i in range(n) for j in range(n)})
    layout.update({f'yaxis{i*n+j+1}': f'y{i*n+j+1}' for i in range(n) for j in range(n)})
    grid_ref = []
    count = 1
    for i in range(n):
        row = []
        for j in range(n):
            row.append([
                MockSubplotRef([f'xaxis{count}']),
                MockSubplotRef([f'yaxis{count}'])
            ])
            count += 1
        grid_ref.append(row)
    fig = MockFigure(grid_ref, layout)
    # Check a few random secondary_y accesses
    codeflash_output = _get_grid_subplot(fig, 1, 1, secondary_y=True)
    codeflash_output = _get_grid_subplot(fig, n, n, secondary_y=True)
    codeflash_output = _get_grid_subplot(fig, n//2, n//2, secondary_y=True)

def test_large_grid_missing_cells():
    # 30x30 grid, some cells empty
    n = 30
    layout = MockLayout({f'xaxis{i*n+j+1}': f'x{i*n+j+1}' for i in range(n) for j in range(n)})
    grid_ref = []
    count = 1
    for i in range(n):
        row = []
        for j in range(n):
            # Make every 10th cell empty
            if (i*n + j) % 10 == 0:
                row.append([])
            else:
                row.append([MockSubplotRef([f'xaxis{count}'])])
            count += 1
        grid_ref.append(row)
    fig = MockFigure(grid_ref, layout)
    # Empty cell
    codeflash_output = _get_grid_subplot(fig, 1, 1)
    # Non-empty cell
    codeflash_output = _get_grid_subplot(fig, 1, 2)
    # Last cell (should not be empty)
    codeflash_output = _get_grid_subplot(fig, n, n)

def test_large_grid_with_domains():
    # 10x10 grid, all with domain only
    n = 10
    grid_ref = []
    for i in range(n):
        row = []
        for j in range(n):
            domain = {'x': [i/n, (i+1)/n], 'y': [j/n, (j+1)/n]}
            row.append([MockSubplotRef(layout_keys=[], trace_kwargs={'domain': domain})])
        grid_ref.append(row)
    layout = MockLayout({})
    fig = MockFigure(grid_ref, layout)
    # Check a few domains
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output
    codeflash_output = _get_grid_subplot(fig, n, n); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import collections

# imports
import pytest  # used for our unit tests
from plotly._subplots import _get_grid_subplot

# function to test
# Constants
# ---------
# Subplot types that are each individually positioned with a domain
#
# Each of these subplot types has a `domain` property with `x`/`y`
# properties.
# Note that this set does not contain `xaxis`/`yaxis` because these behave a
# little differently.

# Named tuple to hold an xaxis/yaxis pair that represent a single subplot
SubplotXY = collections.namedtuple("SubplotXY", ("xaxis", "yaxis"))
SubplotDomain = collections.namedtuple("SubplotDomain", ("x", "y"))
from plotly._subplots import _get_grid_subplot

# Helper classes to simulate fig and subplot_refs structures

class DummySubplotRef:
    """Simulates a subplot reference object with layout_keys and trace_kwargs."""
    def __init__(self, layout_keys, trace_kwargs=None):
        self.layout_keys = layout_keys
        self.trace_kwargs = trace_kwargs or {}

class DummyFig:
    """Simulates a Plotly Figure with _grid_ref and layout."""
    def __init__(self, grid_ref, layout):
        self._grid_ref = grid_ref
        self.layout = layout

# ----------------------------
# Unit Tests for _get_grid_subplot
# ----------------------------

# 1. BASIC TEST CASES

def test_basic_single_subplot_domain():
    # Test when layout_keys is empty, should return SubplotDomain
    domain = {"x": [0, 0.5], "y": [0, 0.5]}
    subplot_ref = DummySubplotRef(layout_keys=[], trace_kwargs={"domain": domain})
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output

def test_basic_single_subplot_layout_key():
    # Test when layout_keys has one key, should return layout[key]
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1"])
    layout = {"xaxis1": "AXIS1"}
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout=layout)
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output

def test_basic_single_subplot_xy():
    # Test when layout_keys has two keys, should return SubplotXY
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1", "yaxis1"])
    layout = {"xaxis1": "AXIS1", "yaxis1": "YAXIS1"}
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout=layout)
    codeflash_output = _get_grid_subplot(fig, 1, 1); result = codeflash_output

def test_basic_secondary_y_returns_second_subplot():
    # Test secondary_y True returns second subplot if present
    subplot_ref1 = DummySubplotRef(layout_keys=["xaxis1", "yaxis1"])
    subplot_ref2 = DummySubplotRef(layout_keys=["xaxis2", "yaxis2"])
    layout = {"xaxis2": "AXIS2", "yaxis2": "YAXIS2"}
    fig = DummyFig(grid_ref=[[ [subplot_ref1, subplot_ref2] ]], layout=layout)
    codeflash_output = _get_grid_subplot(fig, 1, 1, secondary_y=True); result = codeflash_output

def test_basic_secondary_y_returns_none_if_not_present():
    # Test secondary_y True returns None if only one subplot_ref
    subplot_ref1 = DummySubplotRef(layout_keys=["xaxis1", "yaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref1] ]], layout={})
    codeflash_output = _get_grid_subplot(fig, 1, 1, secondary_y=True); result = codeflash_output

# 2. EDGE TEST CASES

def test_missing_grid_ref_attribute():
    # Test fig without _grid_ref raises Exception
    class NoGridFig:
        pass
    fig = NoGridFig()
    with pytest.raises(Exception) as excinfo:
        _get_grid_subplot(fig, 1, 1)

def test_row_out_of_bounds_low():
    # row < 1
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(ValueError) as excinfo:
        _get_grid_subplot(fig, 0, 1)

def test_row_out_of_bounds_high():
    # row > number of rows
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(ValueError) as excinfo:
        _get_grid_subplot(fig, 2, 1)

def test_col_out_of_bounds_low():
    # col < 1
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(ValueError) as excinfo:
        _get_grid_subplot(fig, 1, 0)

def test_col_out_of_bounds_high():
    # col > number of cols
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(ValueError) as excinfo:
        _get_grid_subplot(fig, 1, 2)

def test_row_col_not_integer():
    # row/col not integer
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(ValueError):
        _get_grid_subplot(fig, "1", 1)
    with pytest.raises(ValueError):
        _get_grid_subplot(fig, 1, "1")

def test_empty_subplot_ref_returns_none():
    # grid_ref cell is empty list, should return None
    fig = DummyFig(grid_ref=[[ [] ]], layout={})
    codeflash_output = _get_grid_subplot(fig, 1, 1)

def test_secondary_y_with_one_subplot_ref_returns_none():
    # secondary_y True, but only one subplot_ref
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1", "yaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    codeflash_output = _get_grid_subplot(fig, 1, 1, secondary_y=True)

def test_unexpected_layout_keys_length_raises():
    # layout_keys has 3+ entries, should raise ValueError
    subplot_ref = DummySubplotRef(layout_keys=["xaxis1", "yaxis1", "zaxis1"])
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(ValueError) as excinfo:
        _get_grid_subplot(fig, 1, 1)

def test_domain_missing_keys():
    # layout_keys empty, but domain missing keys, should raise KeyError
    subplot_ref = DummySubplotRef(layout_keys=[], trace_kwargs={"domain": {"x": [0, 1]}})
    fig = DummyFig(grid_ref=[[ [subplot_ref] ]], layout={})
    with pytest.raises(TypeError):
        _get_grid_subplot(fig, 1, 1)

# 3. LARGE SCALE TEST CASES

def test_large_grid_access_first_and_last():
    # Test a large grid (100x100), access first and last cell
    N = 100
    domain = {"x": [0, 1], "y": [0, 1]}
    subplot_ref = DummySubplotRef(layout_keys=[], trace_kwargs={"domain": domain})
    grid_ref = [ [ [subplot_ref] for _ in range(N) ] for _ in range(N) ]
    fig = DummyFig(grid_ref=grid_ref, layout={})
    # First cell
    codeflash_output = _get_grid_subplot(fig, 1, 1); result_first = codeflash_output
    # Last cell
    codeflash_output = _get_grid_subplot(fig, N, N); result_last = codeflash_output

def test_large_grid_secondary_y():
    # Test a large grid (50x50) with secondary_y, all cells have two subplot_refs
    N = 50
    subplot_ref1 = DummySubplotRef(layout_keys=["xaxis1", "yaxis1"])
    subplot_ref2 = DummySubplotRef(layout_keys=["xaxis2", "yaxis2"])
    layout = {"xaxis2": "AXIS2", "yaxis2": "YAXIS2"}
    grid_ref = [ [ [subplot_ref1, subplot_ref2] for _ in range(N) ] for _ in range(N) ]
    fig = DummyFig(grid_ref=grid_ref, layout=layout)
    # Pick a cell in the middle
    codeflash_output = _get_grid_subplot(fig, N//2, N//2, secondary_y=True); result = codeflash_output

def test_large_grid_sparse_empty_cells():
    # Large grid with some empty cells, should return None for empty cell
    N = 20
    domain = {"x": [0, 1], "y": [0, 1]}
    subplot_ref = DummySubplotRef(layout_keys=[], trace_kwargs={"domain": domain})
    grid_ref = []
    for i in range(N):
        row = []
        for j in range(N):
            if (i + j) % 10 == 0:
                row.append([])  # empty cell
            else:
                row.append([subplot_ref])
        grid_ref.append(row)
    fig = DummyFig(grid_ref=grid_ref, layout={})
    # Empty cell
    codeflash_output = _get_grid_subplot(fig, 1, 10)
    # Non-empty cell
    codeflash_output = _get_grid_subplot(fig, 1, 2); result = codeflash_output

def test_large_grid_with_varied_layout_keys():
    # Large grid with varied layout_keys (0, 1, 2), test each type
    N = 10
    domain = {"x": [0, 1], "y": [0, 1]}
    subplot_ref0 = DummySubplotRef(layout_keys=[], trace_kwargs={"domain": domain})
    subplot_ref1 = DummySubplotRef(layout_keys=["xaxis1"])
    subplot_ref2 = DummySubplotRef(layout_keys=["xaxis1", "yaxis1"])
    layout = {"xaxis1": "AXIS1", "yaxis1": "YAXIS1"}
    grid_ref = []
    for i in range(N):
        row = []
        for j in range(N):
            if (i + j) % 3 == 0:
                row.append([subplot_ref0])
            elif (i + j) % 3 == 1:
                row.append([subplot_ref1])
            else:
                row.append([subplot_ref2])
        grid_ref.append(row)
    fig = DummyFig(grid_ref=grid_ref, layout=layout)
    codeflash_output = _get_grid_subplot(fig, 1, 2)
    codeflash_output = _get_grid_subplot(fig, 1, 3); result = codeflash_output
# 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-_get_grid_subplot-mb2hva2a and push.