Skip to content

Conversation

@codeflash-ai
Copy link

@codeflash-ai codeflash-ai bot commented May 24, 2025

📄 110% (1.10x) speedup for get_grid_style in plotly/matplotlylib/mplexporter/utils.py

⏱️ Runtime : 3.98 milliseconds 1.90 milliseconds (best of 417 runs)

📝 Explanation and details

Here's an optimized version of your Python program, targeting your main bottleneck: export_color. From your line profile, most of the runtime is dominated by repeated calls to colorConverter.to_rgba and colorConverter.to_rgb. These functions are pure (for a specific color input always outputting the same results), so caching their outputs aggressively will drastically reduce redundant computations, especially since your usage pattern suggests strong color locality (repeated calls with the same argument).

Other minor tweaks.

  • Replace ",".join(map(str, ...)) in get_dasharray with a faster list comprehension.
  • Avoid repeated dictionary lookups for _dashSeq and in function calls (minor but measurable).
  • In export_color, avoid calling both to_rgba and to_rgb for alpha==1; just slice the values from the first result, keeping only one converter call (since to_rgba is a superset of to_rgb).
  • Replace np.round(val * 255) with int(round(val * 255)) and avoid extra str conversion steps.
  • Use tuple unpacking for string formatting instead of generator expressions in format calls.
  • Use fast-paths and local variable caching.
  • Avoid multiple attribute lookups where possible.

Here is the rewritten, optimized code.

Summary of major optimizations:

  • @lru_cache decorators for color conversions.
  • Avoided duplicate color conversions.
  • More efficient joining and formatting.
  • Reduced repeated dictionary and attribute access.

This rewrite should provide a significant speed up—especially in export_color, cascading into get_grid_style, which the profiler showed as the main hot spot.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 1540 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 🔘 None Found
📊 Tests Coverage 100.0%
🌀 Generated Regression Tests Details
import warnings

import numpy as np
# imports
import pytest  # used for our unit tests
from matplotlib.colors import colorConverter
from plotly.matplotlylib.mplexporter.utils import get_grid_style

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

# Helper mock classes to simulate matplotlib axis and gridline objects

class MockGridLine:
    """Mock object simulating a matplotlib gridline."""
    def __init__(self, color='black', alpha=1.0, linestyle='-', dashseq=None):
        self._color = color
        self._alpha = alpha
        self._linestyle = linestyle
        if dashseq is not None:
            self._dashSeq = dashseq

    def get_color(self):
        return self._color

    def get_alpha(self):
        return self._alpha

    def get_linestyle(self):
        return self._linestyle

class MockAxis:
    """Mock object simulating a matplotlib axis."""
    def __init__(self, gridlines=None, gridOn=True):
        self._gridlines = gridlines if gridlines is not None else []
        self._major_tick_kw = {"gridOn": gridOn}

    def get_gridlines(self):
        return self._gridlines

    @property
    def _major_tick_kw(self):
        return self.__dict__['_major_tick_kw']

    @_major_tick_kw.setter
    def _major_tick_kw(self, value):
        self.__dict__['_major_tick_kw'] = value

# -------- BASIC TEST CASES --------

def test_basic_solid_black_grid():
    """Test basic case: grid is on, one solid black gridline."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_basic_dashed_red_grid():
    """Test grid is on, one dashed red gridline."""
    gridline = MockGridLine(color='red', alpha=0.5, linestyle='--')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_basic_dotted_blue_grid():
    """Test grid is on, one dotted blue gridline."""
    gridline = MockGridLine(color='blue', alpha=0.8, linestyle=':')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_basic_dashdot_green_grid():
    """Test grid is on, one dashdot green gridline."""
    gridline = MockGridLine(color='green', alpha=1.0, linestyle='-.')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

# -------- EDGE TEST CASES --------

def test_grid_off_returns_gridOn_false():
    """Test when gridOn is False, should always return {'gridOn': False}."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=False)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_no_gridlines_returns_gridOn_false():
    """Test when gridOn is True but no gridlines, should return {'gridOn': False}."""
    axis = MockAxis([], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_transparent_gridline_color_none():
    """Test gridline with color=None (fully transparent)."""
    gridline = MockGridLine(color=None, alpha=1.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_alpha_zero():
    """Test gridline with alpha=0 (fully transparent)."""
    gridline = MockGridLine(color='black', alpha=0.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_rgba_color():
    """Test gridline with RGBA color with partial transparency."""
    gridline = MockGridLine(color=(1, 0.5, 0, 0.6), alpha=0.6, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_unknown_linestyle_warns_and_defaults():
    """Test gridline with unknown linestyle, should warn and default to solid."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle='fancy')
    axis = MockAxis([gridline], gridOn=True)
    with pytest.warns(UserWarning):
        codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_dashseq_overrides_linestyle():
    """Test _dashSeq is used instead of get_linestyle()."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle='-', dashseq=[1, 2, 3])
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_blank_linestyle():
    """Test blank linestyle disables the gridline (dasharray None)."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle=' ')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_none_linestyle():
    """Test 'None' linestyle disables the gridline (dasharray None)."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle='None')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_tuple_none_linestyle():
    """Test (None, None) linestyle disables the gridline (dasharray 'none')."""
    gridline = MockGridLine(color='black', alpha=1.0, linestyle=(None, None))
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_color_string_hex():
    """Test gridline with color as hex string."""
    gridline = MockGridLine(color='#FF00FF', alpha=0.7, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_color_tuple_rgb():
    """Test gridline with color as RGB tuple."""
    gridline = MockGridLine(color=(0.5, 0.25, 0.75), alpha=1.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_color_tuple_rgba_alpha1():
    """Test gridline with color as RGBA tuple with alpha=1."""
    gridline = MockGridLine(color=(0.2, 0.4, 0.6, 1.0), alpha=1.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_gridline_with_color_tuple_rgba_alpha0():
    """Test gridline with color as RGBA tuple with alpha=0 (fully transparent)."""
    gridline = MockGridLine(color=(0.2, 0.4, 0.6, 0.0), alpha=0.0, linestyle='-')
    axis = MockAxis([gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

# -------- LARGE SCALE TEST CASES --------

def test_many_gridlines_only_first_used():
    """Test that only the first gridline's style is used, even if many gridlines exist."""
    gridlines = [
        MockGridLine(color='red', alpha=0.2, linestyle='--'),
        MockGridLine(color='blue', alpha=0.8, linestyle=':'),
        MockGridLine(color='green', alpha=1.0, linestyle='-.'),
    ]
    axis = MockAxis(gridlines, gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_large_number_of_gridlines_performance():
    """Test performance and correctness with 1000 gridlines."""
    gridlines = [MockGridLine(color='black', alpha=1.0, linestyle='-') for _ in range(1000)]
    axis = MockAxis(gridlines, gridOn=True)
    codeflash_output = get_grid_style(axis); result = codeflash_output

def test_large_number_of_axes_instances():
    """Test correctness with 1000 axes, each with one gridline."""
    for i in range(1000):
        color = (i % 2 and 'black') or 'blue'
        linestyle = '-' if i % 3 == 0 else '--'
        alpha = 0.5 if i % 4 == 0 else 1.0
        gridline = MockGridLine(color=color, alpha=alpha, linestyle=linestyle)
        axis = MockAxis([gridline], gridOn=True)
        codeflash_output = get_grid_style(axis); result = codeflash_output
        # Check that the result matches expected color and dasharray
        if color == 'black':
            pass
        else:
            pass
        if linestyle == '-':
            pass
        else:
            pass
        if alpha == 0.5:
            pass
        else:
            pass

def test_large_scale_varied_gridlines():
    """Test 500 gridlines with varied colors, linestyles, and alphas."""
    colors = ['red', 'green', 'blue', 'black', (0.5, 0.5, 0.5)]
    linestyles = ['-', '--', ':', '-.', 'None']
    alphas = [0.1, 0.5, 0.9, 1.0]
    for i in range(500):
        color = colors[i % len(colors)]
        linestyle = linestyles[i % len(linestyles)]
        alpha = alphas[i % len(alphas)]
        gridline = MockGridLine(color=color, alpha=alpha, linestyle=linestyle)
        axis = MockAxis([gridline], gridOn=True)
        codeflash_output = get_grid_style(axis); result = codeflash_output
        # Check dasharray matches expected
        if linestyle == '-':
            pass
        elif linestyle == '--':
            pass
        elif linestyle == ':':
            pass
        elif linestyle == '-.':
            pass
        elif linestyle == 'None':
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import warnings

import numpy as np
# imports
import pytest  # used for our unit tests
from matplotlib.colors import colorConverter
from plotly.matplotlylib.mplexporter.utils import get_grid_style

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

# Helper mock classes to simulate axis and gridline objects
class MockGridline:
    def __init__(self, color="black", alpha=1.0, linestyle="solid", dashSeq=None):
        self._color = color
        self._alpha = alpha
        self._linestyle = linestyle
        if dashSeq is not None:
            self._dashSeq = dashSeq

    def get_color(self):
        return self._color

    def get_alpha(self):
        return self._alpha

    def get_linestyle(self):
        return self._linestyle

class MockAxis:
    def __init__(self, gridlines=None, gridOn=True):
        self._gridlines = gridlines if gridlines is not None else []
        self._major_tick_kw = {"gridOn": gridOn}

    def get_gridlines(self):
        return self._gridlines

    @property
    def _major_tick_kw(self):
        return self.__major_tick_kw

    @_major_tick_kw.setter
    def _major_tick_kw(self, value):
        self.__major_tick_kw = value

# 1. BASIC TEST CASES

def test_basic_solid_black_grid():
    # Solid black gridline, gridOn True
    gridline = MockGridline(color="black", alpha=1.0, linestyle="solid")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_basic_dashed_red_grid():
    # Dashed red gridline, gridOn True
    gridline = MockGridline(color="red", alpha=0.7, linestyle="--")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_basic_dotted_blue_grid():
    # Dotted blue gridline, gridOn True
    gridline = MockGridline(color="blue", alpha=0.5, linestyle=":")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_basic_dashdot_green_grid():
    # Dashdot green gridline, gridOn True
    gridline = MockGridline(color="green", alpha=1.0, linestyle="-.")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_basic_grid_off():
    # gridOn False, should return only {"gridOn": False}
    gridline = MockGridline()
    axis = MockAxis(gridlines=[gridline], gridOn=False)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_basic_no_gridlines():
    # gridOn True, but no gridlines present
    axis = MockAxis(gridlines=[], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

# 2. EDGE TEST CASES

def test_edge_alpha_zero():
    # Alpha zero, should still return gridOn True and correct color
    gridline = MockGridline(color="black", alpha=0.0, linestyle="solid")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_none_color():
    # Color is None, should return "none" for color
    gridline = MockGridline(color=None, alpha=1.0, linestyle="solid")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_rgba_color_with_alpha():
    # RGBA color with partial transparency
    gridline = MockGridline(color=(1, 0, 0, 0.5), alpha=0.5, linestyle="solid")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_unknown_linestyle_warns(monkeypatch):
    # Unknown linestyle should warn and default to "none"
    gridline = MockGridline(color="black", alpha=1.0, linestyle="fancy")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    # Patch warnings.warn to capture the warning
    warnings_list = []
    def fake_warn(msg, *args, **kwargs):
        warnings_list.append(msg)
    monkeypatch.setattr(warnings, "warn", fake_warn)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_dashSeq_overrides_linestyle():
    # If _dashSeq is present, use it for dasharray
    gridline = MockGridline(color="black", alpha=1.0, linestyle="solid", dashSeq=[1,2,3])
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_empty_linestyle_maps_to_None():
    # Linestyle "" should map to None dasharray
    gridline = MockGridline(color="black", alpha=1.0, linestyle="")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_space_linestyle_maps_to_None():
    # Linestyle " " should map to None dasharray
    gridline = MockGridline(color="black", alpha=1.0, linestyle=" ")
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_none_linestyle_maps_to_None():
    # Linestyle None should map to None dasharray
    gridline = MockGridline(color="black", alpha=1.0, linestyle=None)
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_multiple_gridlines_only_first_used():
    # Multiple gridlines, only first is used
    gridline1 = MockGridline(color="red", alpha=0.3, linestyle="--")
    gridline2 = MockGridline(color="blue", alpha=0.7, linestyle=":")
    axis = MockAxis(gridlines=[gridline1, gridline2], gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_edge_major_tick_kw_missing_gridOn(monkeypatch):
    # _major_tick_kw missing "gridOn" key should raise KeyError
    gridline = MockGridline()
    axis = MockAxis(gridlines=[gridline], gridOn=True)
    # Remove "gridOn" key
    axis._major_tick_kw = {}
    with pytest.raises(KeyError):
        get_grid_style(axis)

# 3. LARGE SCALE TEST CASES

def test_large_many_gridlines():
    # Many gridlines, only first is used, should be fast
    gridlines = [MockGridline(color="black", alpha=1.0, linestyle="solid") for _ in range(500)]
    axis = MockAxis(gridlines=gridlines, gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_large_varied_gridlines():
    # Many gridlines with varied properties, only first is used
    gridlines = []
    for i in range(500):
        color = "red" if i % 2 == 0 else "blue"
        alpha = 1.0 if i % 3 == 0 else 0.5
        linestyle = "--" if i % 4 == 0 else ":"
        gridlines.append(MockGridline(color=color, alpha=alpha, linestyle=linestyle))
    axis = MockAxis(gridlines=gridlines, gridOn=True)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_large_gridlines_gridOff():
    # Many gridlines, but gridOn is False
    gridlines = [MockGridline(color="black", alpha=1.0, linestyle="solid") for _ in range(500)]
    axis = MockAxis(gridlines=gridlines, gridOn=False)
    codeflash_output = get_grid_style(axis); style = codeflash_output

def test_large_no_gridlines_gridOn():
    # gridOn True but no gridlines
    axis = MockAxis(gridlines=[], gridOn=True)
    codeflash_output = get_grid_style(axis); style = 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_style-mb27tu6y and push.

Codeflash

Here's an optimized version of your Python program, targeting your main bottleneck: **export_color**. From your line profile, most of the runtime is dominated by repeated calls to `colorConverter.to_rgba` and `colorConverter.to_rgb`. These functions are pure (for a specific color input always outputting the same results), so **caching** their outputs aggressively will drastically reduce redundant computations, especially since your usage pattern suggests strong color locality (repeated calls with the same argument).

Other minor tweaks.

- Replace `",".join(map(str, ...))` in `get_dasharray` with a faster list comprehension.
- Avoid repeated dictionary lookups for `_dashSeq` and in function calls (minor but measurable).
- In `export_color`, avoid calling both `to_rgba` and `to_rgb` for alpha==1; just slice the values from the first result, keeping only one converter call (since `to_rgba` is a superset of `to_rgb`).
- Replace `np.round(val * 255)` with `int(round(val * 255))` and avoid extra `str` conversion steps.
- Use tuple unpacking for string formatting instead of generator expressions in format calls.
- Use fast-paths and local variable caching.
- Avoid multiple attribute lookups where possible.

Here is the rewritten, optimized code.



**Summary of major optimizations:**

- **`@lru_cache`** decorators for color conversions.
- **Avoided duplicate color conversions.**
- **More efficient joining and formatting.**
- **Reduced repeated dictionary and attribute access.**

This rewrite should provide a significant speed up—especially in `export_color`, cascading into `get_grid_style`, which the profiler showed as the main hot spot.
@codeflash-ai codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label May 24, 2025
@codeflash-ai codeflash-ai bot requested a review from misrasaurabh1 May 24, 2025 12:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

⚡️ codeflash Optimization PR opened by Codeflash AI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants