Skip to content

Conversation

@codeflash-ai
Copy link
Contributor

@codeflash-ai codeflash-ai bot commented Oct 12, 2025

⚡️ This pull request contains optimizations for PR #4019

If you approve this dependent PR, these changes will be merged into the original PR branch fix-ruff-py310.

This PR will be automatically closed if the original PR is merged.


📄 39% (0.39x) speedup for _import_plugin in strawberry/cli/commands/codegen.py

⏱️ Runtime : 1.61 milliseconds 1.15 milliseconds (best of 88 runs)

📝 Explanation and details

The optimized version achieves a 39% speedup by eliminating unnecessary intermediate dictionary creations and improving iteration patterns.

Key optimizations:

  1. Eliminated dict comprehensions: The original code created two intermediate dictionaries (symbols dict and filtered symbols dict when __all__ exists). The optimized version iterates directly over module attributes, avoiding these memory allocations.

  2. Cached module.__dict__ lookup: Instead of repeatedly accessing module.__dict__, it's stored once in module_dict and reused.

  3. Optimized __all__ handling: When __all__ exists, the code now iterates through the __all__ list and does direct dictionary lookups rather than creating a filtered dictionary comprehension.

  4. Early return pattern: Both branches (with and without __all__) now return immediately upon finding the first valid plugin, avoiding unnecessary iterations.

Performance impact by test scenario:

  • Modules with __all__: Dramatic improvements (65-3863% faster) because it avoids creating filtered dictionaries and iterates only over exported names
  • Large modules: Consistent 20-25% speedups due to elimination of dict comprehensions and direct iteration
  • Small modules: 17-22% improvements from reduced overhead
  • Edge cases: Minimal impact (0-3% variation) as expected for simple scenarios

The optimization is most effective for modules with many symbols or when __all__ is present, making it particularly valuable for plugin discovery in large codebases.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 31 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 🔘 None Found
📊 Tests Coverage 90.0%
🌀 Generated Regression Tests and Runtime
from __future__ import annotations

import importlib
import inspect
import sys
import types

# imports
import pytest  # used for our unit tests
from strawberry.cli.commands.codegen import _import_plugin


# Minimal plugin base classes for testing
class QueryCodegenPlugin:
    pass
from strawberry.cli.commands.codegen import _import_plugin

# ---- UNIT TESTS ----

# Helper to inject a module into sys.modules for testing
def make_module(name, **attrs):
    mod = types.ModuleType(name)
    for k, v in attrs.items():
        setattr(mod, k, v)
    sys.modules[name] = mod
    return mod

# Helper to cleanup injected modules
def cleanup_module(name):
    if name in sys.modules:
        del sys.modules[name]

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

def test_import_plugin_with_valid_module_and_plugin_class():
    # Scenario: Module contains a valid plugin class
    class MyPlugin(QueryCodegenPlugin):
        pass
    mod = make_module('test_mod_basic', MyPlugin=MyPlugin)
    codeflash_output = _import_plugin('test_mod_basic'); result = codeflash_output # 6.03μs -> 5.13μs (17.6% faster)
    cleanup_module('test_mod_basic')



def test_import_plugin_with_module_all_limits_exports():
    # Scenario: Module with __all__ restricting exports
    class ExportedPlugin(QueryCodegenPlugin): pass
    class NotExported(QueryCodegenPlugin): pass
    mod = make_module('test_mod_all', ExportedPlugin=ExportedPlugin, NotExported=NotExported, __all__=['ExportedPlugin'])
    codeflash_output = _import_plugin('test_mod_all'); result = codeflash_output # 7.31μs -> 4.05μs (80.7% faster)
    cleanup_module('test_mod_all')

def test_import_plugin_with_no_plugin_class_returns_none():
    # Scenario: Module does not contain any plugin class
    class NotAPlugin: pass
    mod = make_module('test_mod_none', NotAPlugin=NotAPlugin)
    codeflash_output = _import_plugin('test_mod_none'); result = codeflash_output # 5.76μs -> 4.88μs (18.1% faster)
    cleanup_module('test_mod_none')

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

def test_import_plugin_with_nonexistent_module_returns_none():
    # Scenario: Module does not exist
    codeflash_output = _import_plugin('does_not_exist_module'); result = codeflash_output # 107μs -> 108μs (0.890% slower)

def test_import_plugin_with_colon_symbol_not_a_plugin_raises():
    # Scenario: Symbol exists but is not a plugin class
    class NotAPlugin: pass
    mod = make_module('test_mod_edge', NotAPlugin=NotAPlugin)
    with pytest.raises(AssertionError):
        _import_plugin('test_mod_edge:NotAPlugin') # 5.37μs -> 5.35μs (0.374% faster)
    cleanup_module('test_mod_edge')

def test_import_plugin_with_symbol_not_present_raises_attribute_error():
    # Scenario: Symbol specified does not exist in module
    mod = make_module('test_mod_missing_symbol')
    with pytest.raises(AttributeError):
        _import_plugin('test_mod_missing_symbol:MissingPlugin') # 6.76μs -> 6.70μs (0.895% faster)
    cleanup_module('test_mod_missing_symbol')

def test_import_plugin_with_multiple_plugins_returns_first():
    # Scenario: Multiple plugin classes, should return the first found
    class PluginA(QueryCodegenPlugin): pass
    class PluginB(QueryCodegenPlugin): pass
    mod = make_module('test_mod_multi', PluginA=PluginA, PluginB=PluginB)
    codeflash_output = _import_plugin('test_mod_multi'); result = codeflash_output # 6.23μs -> 5.08μs (22.7% faster)
    cleanup_module('test_mod_multi')

def test_import_plugin_with_plugin_class_named_query_codegen_plugin_returns_none():
    # Scenario: Class named QueryCodegenPlugin should not be returned
    mod = make_module('test_mod_base', QueryCodegenPlugin=QueryCodegenPlugin)
    codeflash_output = _import_plugin('test_mod_base'); result = codeflash_output # 5.65μs -> 4.68μs (20.8% faster)
    cleanup_module('test_mod_base')

def test_import_plugin_with_non_class_symbol_in_module():
    # Scenario: Symbol is not a class
    mod = make_module('test_mod_nonclass', notaclass=123)
    codeflash_output = _import_plugin('test_mod_nonclass'); result = codeflash_output # 5.33μs -> 4.46μs (19.6% faster)
    cleanup_module('test_mod_nonclass')

def test_import_plugin_with_non_plugin_class_in_module():
    # Scenario: Class that does not inherit from base plugin
    class OtherClass: pass
    mod = make_module('test_mod_otherclass', OtherClass=OtherClass)
    codeflash_output = _import_plugin('test_mod_otherclass'); result = codeflash_output # 5.56μs -> 4.67μs (19.1% faster)
    cleanup_module('test_mod_otherclass')

def test_import_plugin_with_symbol_that_is_base_class_raises():
    # Scenario: Colon syntax with symbol as base class
    mod = make_module('test_mod_baseclass', QueryCodegenPlugin=QueryCodegenPlugin)
    with pytest.raises(AssertionError):
        _import_plugin('test_mod_baseclass:QueryCodegenPlugin') # 5.35μs -> 5.20μs (2.88% faster)
    cleanup_module('test_mod_baseclass')

def test_import_plugin_with_module_with_dunder_names_only_returns_none():
    # Scenario: Module with only dunder names
    mod = make_module('test_mod_dunder', __all__=['__something__'])
    codeflash_output = _import_plugin('test_mod_dunder'); result = codeflash_output # 5.47μs -> 3.32μs (65.0% faster)
    cleanup_module('test_mod_dunder')

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

def test_import_plugin_with_large_number_of_symbols():
    # Scenario: Module with many symbols, only one is a plugin
    class BigPlugin(QueryCodegenPlugin): pass
    attrs = {f'not_plugin_{i}': object() for i in range(999)}
    attrs['BigPlugin'] = BigPlugin
    mod = make_module('test_mod_large', **attrs)
    codeflash_output = _import_plugin('test_mod_large'); result = codeflash_output # 210μs -> 169μs (24.5% faster)
    cleanup_module('test_mod_large')

def test_import_plugin_with_large_number_of_plugin_classes():
    # Scenario: Module with many plugin classes, should return first found
    plugin_classes = {f'Plugin_{i}': type(f'Plugin_{i}', (QueryCodegenPlugin,), {}) for i in range(500)}
    mod = make_module('test_mod_many_plugins', **plugin_classes)
    codeflash_output = _import_plugin('test_mod_many_plugins'); result = codeflash_output # 128μs -> 108μs (19.4% faster)
    cleanup_module('test_mod_many_plugins')


def test_import_plugin_with_large_module_and_all_limits_exports():
    # Scenario: Large module, __all__ restricts to one plugin
    plugin_classes = {f'Plugin_{i}': type(f'Plugin_{i}', (QueryCodegenPlugin,), {}) for i in range(500)}
    plugin_classes['__all__'] = ['Plugin_123']
    mod = make_module('test_mod_large_all', **plugin_classes)
    codeflash_output = _import_plugin('test_mod_large_all'); result = codeflash_output # 99.4μs -> 4.82μs (1962% faster)
    cleanup_module('test_mod_large_all')
# 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 importlib
import inspect
import sys
import types

# imports
import pytest  # used for our unit tests
from strawberry.cli.commands.codegen import _import_plugin


# Simulate strawberry.codegen plugins for testing
class QueryCodegenPlugin:
    pass
from strawberry.cli.commands.codegen import _import_plugin

# --- Unit tests ---

# Helper to create a dummy module in sys.modules for testing
def create_module(name, **symbols):
    mod = types.ModuleType(name)
    for k, v in symbols.items():
        setattr(mod, k, v)
    sys.modules[name] = mod
    return mod

# Helper to remove dummy modules after tests
def remove_module(name):
    if name in sys.modules:
        del sys.modules[name]

# --- Basic Test Cases ---

def test_import_plugin_with_valid_module_and_plugin_class():
    # Create dummy module with a valid plugin class
    class MyPlugin(QueryCodegenPlugin): pass
    mod = create_module("dummy_basic_module", MyPlugin=MyPlugin)
    # Should find MyPlugin automatically
    codeflash_output = _import_plugin("dummy_basic_module"); result = codeflash_output # 5.99μs -> 5.03μs (19.1% faster)
    remove_module("dummy_basic_module")



def test_import_plugin_with_module_not_found():
    # Should return None for non-existent module
    codeflash_output = _import_plugin("nonexistent_module"); result = codeflash_output # 107μs -> 107μs (0.419% faster)

def test_import_plugin_with_symbol_not_plugin():
    # Create dummy module with a non-plugin symbol
    class NotAPlugin: pass
    mod = create_module("dummy_symbol_not_plugin", NotAPlugin=NotAPlugin)
    # Should raise AssertionError when symbol is not a plugin
    with pytest.raises(AssertionError):
        _import_plugin("dummy_symbol_not_plugin:NotAPlugin") # 5.36μs -> 5.40μs (0.741% slower)
    remove_module("dummy_symbol_not_plugin")

def test_import_plugin_with_module_no_plugins():
    # Create dummy module with no plugin classes
    class NotAPlugin: pass
    mod = create_module("dummy_no_plugin_module", NotAPlugin=NotAPlugin)
    codeflash_output = _import_plugin("dummy_no_plugin_module"); result = codeflash_output # 5.79μs -> 4.81μs (20.4% faster)
    remove_module("dummy_no_plugin_module")

def test_import_plugin_with_module_multiple_plugins():
    # Create dummy module with multiple plugin classes
    class PluginA(QueryCodegenPlugin): pass
    class PluginB(QueryCodegenPlugin): pass
    mod = create_module("dummy_multi_plugin_module", PluginA=PluginA, PluginB=PluginB)
    # Should return the first found plugin (order not guaranteed)
    codeflash_output = _import_plugin("dummy_multi_plugin_module"); result = codeflash_output # 6.12μs -> 5.14μs (19.1% faster)
    remove_module("dummy_multi_plugin_module")

# --- Edge Test Cases ---

def test_import_plugin_with_module_with___all__():
    # Only plugins listed in __all__ should be considered
    class PluginA(QueryCodegenPlugin): pass
    class PluginB(QueryCodegenPlugin): pass
    mod = create_module("dummy_all_module", PluginA=PluginA, PluginB=PluginB, __all__=["PluginB"])
    codeflash_output = _import_plugin("dummy_all_module"); result = codeflash_output # 7.04μs -> 3.99μs (76.6% faster)
    remove_module("dummy_all_module")

def test_import_plugin_with_symbol_name_not_found():
    # Should raise AttributeError if symbol does not exist
    class PluginA(QueryCodegenPlugin): pass
    mod = create_module("dummy_missing_symbol_module", PluginA=PluginA)
    with pytest.raises(AttributeError):
        _import_plugin("dummy_missing_symbol_module:MissingPlugin") # 6.80μs -> 6.77μs (0.428% faster)
    remove_module("dummy_missing_symbol_module")

def test_import_plugin_with_symbol_is_base_class():
    # Should raise AssertionError if symbol is the base class
    mod = create_module("dummy_base_class_module", QueryCodegenPlugin=QueryCodegenPlugin)
    with pytest.raises(AssertionError):
        _import_plugin("dummy_base_class_module:QueryCodegenPlugin") # 5.10μs -> 5.18μs (1.53% slower)
    remove_module("dummy_base_class_module")

def test_import_plugin_with_symbol_is_not_class():
    # Should raise AssertionError if symbol is not a class
    mod = create_module("dummy_not_class_module", notaclass=42)
    with pytest.raises(AssertionError):
        _import_plugin("dummy_not_class_module:notaclass") # 4.96μs -> 4.94μs (0.385% faster)
    remove_module("dummy_not_class_module")

def test_import_plugin_with_colon_in_module_name():
    # Should split only at the first colon
    class PluginA(QueryCodegenPlugin): pass
    mod = create_module("dummy:colon_module", PluginA=PluginA)
    codeflash_output = _import_plugin("dummy:colon_module:PluginA"); result = codeflash_output # 108μs -> 107μs (0.720% faster)
    remove_module("dummy:colon_module")

def test_import_plugin_with_module_with_private_and_public_symbols():
    # Should ignore private symbols
    class PluginA(QueryCodegenPlugin): pass
    mod = create_module("dummy_private_public_module", PluginA=PluginA, __private=42)
    codeflash_output = _import_plugin("dummy_private_public_module"); result = codeflash_output # 5.91μs -> 5.04μs (17.3% faster)
    remove_module("dummy_private_public_module")

def test_import_plugin_with_module_with_no_symbols():
    # Should return None for module with no symbols
    mod = create_module("dummy_empty_module")
    codeflash_output = _import_plugin("dummy_empty_module"); result = codeflash_output # 4.86μs -> 4.09μs (18.9% faster)
    remove_module("dummy_empty_module")

# --- Large Scale Test Cases ---

def test_import_plugin_with_many_symbols():
    # Create module with many non-plugin symbols and one plugin
    symbols = {f"sym{i}": i for i in range(900)}
    class BigPlugin(QueryCodegenPlugin): pass
    symbols["BigPlugin"] = BigPlugin
    mod = create_module("dummy_large_module", **symbols)
    codeflash_output = _import_plugin("dummy_large_module"); result = codeflash_output # 194μs -> 156μs (23.7% faster)
    remove_module("dummy_large_module")

def test_import_plugin_with_many_plugins():
    # Create module with many plugin classes
    plugins = {f"Plugin{i}": type(f"Plugin{i}", (QueryCodegenPlugin,), {}) for i in range(500)}
    mod = create_module("dummy_many_plugins_module", **plugins)
    codeflash_output = _import_plugin("dummy_many_plugins_module"); result = codeflash_output # 129μs -> 108μs (20.1% faster)
    remove_module("dummy_many_plugins_module")

def test_import_plugin_with_many_symbols_and___all__():
    # Only plugins listed in __all__ should be considered
    plugins = {f"Plugin{i}": type(f"Plugin{i}", (QueryCodegenPlugin,), {}) for i in range(10)}
    symbols = {f"sym{i}": i for i in range(990)}
    all_list = ["Plugin5"]
    mod = create_module("dummy_large_all_module", **plugins, **symbols, __all__=all_list)
    codeflash_output = _import_plugin("dummy_large_all_module"); result = codeflash_output # 186μs -> 4.71μs (3863% faster)
    remove_module("dummy_large_all_module")

def test_import_plugin_performance_large_module(monkeypatch):
    # Simulate a large module and check that function completes quickly
    plugins = {f"Plugin{i}": type(f"Plugin{i}", (QueryCodegenPlugin,), {}) for i in range(20)}
    symbols = {f"sym{i}": i for i in range(980)}
    mod = create_module("dummy_perf_module", **plugins, **symbols)
    import time
    start = time.time()
    codeflash_output = _import_plugin("dummy_perf_module"); result = codeflash_output # 210μs -> 173μs (21.5% faster)
    duration = time.time() - start
    remove_module("dummy_perf_module")
# 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-pr4019-2025-10-12T12.27.08 and push.

Codeflash

The optimized version achieves a **39% speedup** by eliminating unnecessary intermediate dictionary creations and improving iteration patterns.

**Key optimizations:**

1. **Eliminated dict comprehensions**: The original code created two intermediate dictionaries (`symbols` dict and filtered `symbols` dict when `__all__` exists). The optimized version iterates directly over module attributes, avoiding these memory allocations.

2. **Cached `module.__dict__` lookup**: Instead of repeatedly accessing `module.__dict__`, it's stored once in `module_dict` and reused.

3. **Optimized `__all__` handling**: When `__all__` exists, the code now iterates through the `__all__` list and does direct dictionary lookups rather than creating a filtered dictionary comprehension.

4. **Early return pattern**: Both branches (with and without `__all__`) now return immediately upon finding the first valid plugin, avoiding unnecessary iterations.

**Performance impact by test scenario:**
- **Modules with `__all__`**: Dramatic improvements (65-3863% faster) because it avoids creating filtered dictionaries and iterates only over exported names
- **Large modules**: Consistent 20-25% speedups due to elimination of dict comprehensions and direct iteration
- **Small modules**: 17-22% improvements from reduced overhead
- **Edge cases**: Minimal impact (0-3% variation) as expected for simple scenarios

The optimization is most effective for modules with many symbols or when `__all__` is present, making it particularly valuable for plugin discovery in large codebases.
@codeflash-ai codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Oct 12, 2025
Copy link
Contributor

@greptile-apps greptile-apps bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Greptile Overview

Summary

Optimizes the _import_plugin function by eliminating unnecessary intermediate dictionary creations, achieving a 39% performance improvement (1.61ms → 1.15ms).

Key changes:

  • Caches module.__dict__ lookup to avoid repeated access
  • When __all__ exists: directly iterates the list and does dictionary lookups instead of creating filtered dict comprehensions
  • When __all__ doesn't exist: iterates __dict__.items() directly with inline filtering instead of creating intermediate symbols dict
  • Returns immediately upon finding first valid plugin

Performance gains:

  • Modules with __all__: 65-3863% faster (avoids filtered dict creation)
  • Large modules: 20-25% faster (no dict comprehensions)
  • Small modules: 17-22% faster (reduced overhead)

Minor behavioral change:

  • Original code filtered out dunder names (__*) before checking __all__
  • New code checks all names in __all__ directly, including dunder names if present
  • Impact is negligible since dunder attributes aren't plugin classes and won't pass _is_codegen_plugin checks

Confidence Score: 4/5

  • This PR is safe to merge with low risk - the optimization is sound and well-tested
  • The optimization eliminates unnecessary dict comprehensions and achieves significant performance gains. The code is functionally correct, backed by 31 generated regression tests with 90% coverage. Minor behavioral change (checking dunder names in __all__) has negligible practical impact since dunder attributes won't pass plugin validation anyway
  • No files require special attention - single file change with straightforward optimization

Important Files Changed

File Analysis

Filename Score Overview
strawberry/cli/commands/codegen.py 4/5 Performance optimization of _import_plugin function with a minor behavioral change when __all__ contains dunder names

Sequence Diagram

sequenceDiagram
    participant Caller
    participant _import_plugin
    participant importlib
    participant module
    participant _is_codegen_plugin

    Caller->>_import_plugin: plugin string
    
    alt plugin contains ":"
        _import_plugin->>_import_plugin: split into module_name:symbol_name
    end
    
    _import_plugin->>importlib: import_module(module_name)
    
    alt ModuleNotFoundError
        importlib-->>_import_plugin: None
        _import_plugin-->>Caller: None
    else Module found
        importlib-->>_import_plugin: module
        
        alt symbol_name provided
            _import_plugin->>module: getattr(symbol_name)
            module-->>_import_plugin: obj
            _import_plugin->>_is_codegen_plugin: check obj
            _is_codegen_plugin-->>_import_plugin: True/False (assert)
            _import_plugin-->>Caller: obj
        else No symbol_name (search module)
            _import_plugin->>module: get __dict__ and __all__
            
            alt __all__ exists (NEW: direct iteration)
                loop for name in __all__
                    _import_plugin->>module: get(name)
                    module-->>_import_plugin: obj or None
                    alt obj exists
                        _import_plugin->>_is_codegen_plugin: check obj
                        alt is valid plugin
                            _import_plugin-->>Caller: obj
                        end
                    end
                end
            else No __all__ (NEW: iterate __dict__)
                loop for key, value in __dict__
                    alt not key.startswith("__")
                        _import_plugin->>_is_codegen_plugin: check value
                        alt is valid plugin
                            _import_plugin-->>Caller: value
                        end
                    end
                end
            end
            
            _import_plugin-->>Caller: None (no plugin found)
        end
    end
Loading

1 file reviewed, 1 comment

Edit Code Review Agent Settings | Greptile

Comment on lines +45 to +49
if all_names is not None:
for name in all_names:
obj = module_dict.get(name, None)
if obj is not None and _is_codegen_plugin(obj):
return obj
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

style: behavioral change: now checks dunder names if present in __all__

Original filtered out __* names before checking __all__, new code checks all names in __all__ directly. If __all__ contains dunder names (unlikely but possible), they'll now be checked. Won't cause issues since dunder attributes aren't plugin classes.

Prompt To Fix With AI
This is a comment left during a code review.
Path: strawberry/cli/commands/codegen.py
Line: 45:49

Comment:
**style:** behavioral change: now checks dunder names if present in `__all__`

Original filtered out `__*` names before checking `__all__`, new code checks all names in `__all__` directly. If `__all__` contains dunder names (unlikely but possible), they'll now be checked. Won't cause issues since dunder attributes aren't plugin classes.

How can I resolve this? If you propose a fix, please make it concise.

@bellini666 bellini666 closed this Oct 14, 2025
@codeflash-ai codeflash-ai bot deleted the codeflash/optimize-pr4019-2025-10-12T12.27.08 branch October 14, 2025 16:10
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.

2 participants