⚡️ Speed up function resolve_original_schema by 35% in pydantic/_internal/_generate_schema.py

[codeflash-ai]

📄 resolve_original_schema() in pydantic/_internal/_generate_schema.py

📈 Performance improved by 35% (0.35x faster)

⏱️ Runtime went down from 11.2 microseconds to 8.29 microseconds

Explanation and details

Here's an optimized version of the given Python program.

Changes Made

1. Avoid Multiple Dictionary Lookups: By storing schema['type'] in the variable schema_type at the beginning, we avoid accessing the dictionary multiple times, which can be a relatively slow operation.
2. Simplified Return Statement: Removed the , None part in definitions.get(schema['schema_ref']) because dict.get() already returns None if the key is not found, making this explicit mention redundant and adding a small unnecessary overhead.

These changes ensure that the code runs faster while maintaining the same functionality and return values.

Correctness verification

The new optimized code was tested for correctness. The results are listed below.

🔘 (none found) − ⚙️ Existing Unit Tests

✅ 18 Passed − 🌀 Generated Regression Tests

(click to show generated tests)

# imports
import pytest  # used for our unit tests
from pydantic._internal._generate_schema import resolve_original_schema

# unit tests

def test_definition_ref_resolution():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'string'}

def test_definitions_handling():
    schema = {'type': 'definitions', 'schema': {'type': 'object', 'properties': {'name': {'type': 'string'}}}}
    definitions = {}
    assert resolve_original_schema(schema, definitions) == {'type': 'object', 'properties': {'name': {'type': 'string'}}}

def test_empty_schema():
    schema = {}
    definitions = {}
    assert resolve_original_schema(schema, definitions) == {}

def test_missing_type_key():
    schema = {'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'schema_ref': 'ref1'}

def test_unknown_type_value():
    schema = {'type': 'unknown-type', 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'unknown-type', 'schema_ref': 'ref1'}

def test_missing_schema_ref_key():
    schema = {'type': 'definition-ref'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) is None

def test_empty_definitions_dictionary():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1'}
    definitions = {}
    assert resolve_original_schema(schema, definitions) is None

def test_nested_definitions():
    schema = {'type': 'definitions', 'schema': {'type': 'object', 'properties': {'address': {'type': 'definition-ref', 'schema_ref': 'addressRef'}}}}
    definitions = {'addressRef': {'type': 'object', 'properties': {'street': {'type': 'string'}}}}
    assert resolve_original_schema(schema, definitions) == {'type': 'object', 'properties': {'address': {'type': 'definition-ref', 'schema_ref': 'addressRef'}}}

def test_circular_references():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'definition-ref', 'schema_ref': 'ref2'}, 'ref2': {'type': 'definition-ref', 'schema_ref': 'ref1'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'definition-ref', 'schema_ref': 'ref2'}

def test_large_definitions_dictionary():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref9999'}
    definitions = {f'ref{i}': {'type': 'integer'} for i in range(10000)}
    assert resolve_original_schema(schema, definitions) == {'type': 'integer'}

def test_large_nested_schema():
    schema = {'type': 'definitions', 'schema': {'type': 'object', 'properties': {f'field{i}': {'type': 'string'} for i in range(1000)}}}
    definitions = {}
    assert resolve_original_schema(schema, definitions) == {'type': 'object', 'properties': {f'field{i}': {'type': 'string'} for i in range(1000)}}

def test_consistent_results():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'string'}
    assert resolve_original_schema(schema, definitions) == {'type': 'string'}

def test_schema_with_invalid_type():
    schema = {'type': 123, 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 123, 'schema_ref': 'ref1'}

def test_schema_with_invalid_structure():
    schema = {'type': 'definition-ref', 'schema_ref': 123}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) is None

def test_schema_with_additional_unexpected_keys():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1', 'extra_key': 'extra_value'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'string'}

def test_schema_with_null_values():
    schema = {'type': None, 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': None, 'schema_ref': 'ref1'}

def test_schema_with_mixed_data_types():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1', 'mixed_type': [1, 'string', None]}
    definitions = {'ref1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'string'}

def test_schema_with_special_characters_in_keys():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref@1'}
    definitions = {'ref@1': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'string'}

def test_schema_with_deeply_nested_structures():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'object', 'properties': {'nested': {'type': 'definition-ref', 'schema_ref': 'ref2'}}}, 'ref2': {'type': 'string'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'object', 'properties': {'nested': {'type': 'definition-ref', 'schema_ref': 'ref2'}}}

def test_schema_with_circular_references_in_definitions():
    schema = {'type': 'definition-ref', 'schema_ref': 'ref1'}
    definitions = {'ref1': {'type': 'definition-ref', 'schema_ref': 'ref2'}, 'ref2': {'type': 'definition-ref', 'schema_ref': 'ref1'}}
    assert resolve_original_schema(schema, definitions) == {'type': 'definition-ref', 'schema_ref': 'ref2'}

✅ 1 Passed − ⏪ Replay Tests

[iusedmyimagination]

Looks legit.