⚡️ Speed up function rail_string_to_schema by 61%

[codeflash-ai]

📄 61% (0.61x) speedup for rail_string_to_schema in guardrails/schema/rail_schema.py

⏱️ Runtime : 34.4 milliseconds → 21.4 milliseconds (best of 28 runs)

📝 Explanation and details

The optimized code achieves a 60% speedup by implementing two key optimizations:

1. Pre-cached ValidationType instances (_VALIDATION_TYPE_MAPPING)
The original code creates a new ValidationType object for each schema element, which is expensive. The optimization introduces a module-level cache that pre-creates these instances:

_VALIDATION_TYPE_MAPPING = {
    RailTypes.STRING: ValidationType(SimpleTypes.STRING),
    RailTypes.INTEGER: ValidationType(SimpleTypes.INTEGER),
    # ... etc
}

This eliminates repeated object construction - from the profiler, we can see ValidationType creation time drops from ~28ms to ~0.2ms across all calls.

2. Optimized xml_to_string calls
The original code calls xml_to_string() on every attribute access, even when the attribute is already a string or None. The optimization adds type checks to avoid unnecessary conversions:

# Before: Always calls xml_to_string
description = xml_to_string(element.attrib.get("description"))

# After: Only calls xml_to_string when needed
description_raw = element.attrib.get("description")
description = description_raw if (description_raw is None or isinstance(description_raw, str)) else xml_to_string(description_raw)

Performance Impact by Test Case:

• Large object schemas see the biggest gains (84.8% faster) because they create many ValidationType instances
• Basic schemas see 25-35% improvements from reduced ValidationType construction overhead
• Enum and choice schemas benefit from both optimizations, showing 25-36% speedup
• Edge cases show smaller but consistent improvements (2-7%) since they hit fewer optimized paths

The optimizations are most effective for schemas with many elements that require type validation, which represents the common use case for this XML schema parsing functionality.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
unit_tests/actions/test_reask.py::test_get_reask_prompt	413μs	308μs	33.9%✅
unit_tests/test_rail.py::test_format_not_read_as_validators	221μs	160μs	38.3%✅
unit_tests/test_rail.py::test_rail_list_with_list	244μs	178μs	37.1%✅
unit_tests/test_rail.py::test_rail_list_with_object	253μs	184μs	37.6%✅
unit_tests/test_rail.py::test_rail_list_with_scalar	217μs	159μs	36.3%✅
unit_tests/test_rail.py::test_rail_object_with_list	290μs	209μs	38.9%✅
unit_tests/test_rail.py::test_rail_object_with_object	295μs	211μs	39.6%✅
unit_tests/test_rail.py::test_rail_object_with_scalar	287μs	202μs	42.2%✅
unit_tests/test_rail.py::test_rail_outermost_list	220μs	167μs	32.0%✅
unit_tests/test_rail.py::test_rail_scalar_string	222μs	161μs	38.1%✅
unit_tests/test_rail.py::test_rail_with_messages	223μs	161μs	38.4%✅
unit_tests/test_skeleton.py::test_skeleton	2.02ms	1.41ms	43.1%✅

🌀 Generated Regression Tests and Runtime

from string import Template
# Minimal stubs for required classes and enums to support the function.
from typing import Any, Dict, List, Optional, Union

# imports
import pytest
from guardrails.schema.rail_schema import rail_string_to_schema
# Begin function definition
from lxml import etree as ET
from lxml.etree import XMLParser, _Element


class SimpleTypes:
    STRING = "string"
    INTEGER = "integer"
    NUMBER = "number"
    BOOLEAN = "boolean"
    OBJECT = "object"
    ARRAY = "array"

class OutputTypes:
    STRING = "string"
    LIST = "list"
    DICT = "dict"

class GuardExecutionOptions:
    def __init__(self):
        self.messages = []
        self.reask_messages = []

class ProcessedSchema:
    def __init__(
        self,
        validators: Optional[List[Any]] = None,
        validator_map: Optional[Dict[str, Any]] = None,
        exec_opts: Optional[GuardExecutionOptions] = None,
    ):
        self.validators = validators if validators is not None else []
        self.validator_map = validator_map if validator_map is not None else {}
        self.exec_opts = exec_opts if exec_opts is not None else GuardExecutionOptions()
        self.json_schema = None
        self.output_type = None
from guardrails.schema.rail_schema import rail_string_to_schema

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

# Basic Test Cases

def test_basic_string_output():
    """Test simple string output schema."""
    rail = """<rail>
        <output type="string" description="A simple string." />
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 146μs -> 116μs (26.3% faster)

def test_basic_object_output():
    """Test simple object output schema with one field."""
    rail = """<rail>
        <output type="object" description="An object.">
            <string name="foo" description="A foo string." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 190μs -> 145μs (31.1% faster)

def test_basic_list_output():
    """Test simple list output schema."""
    rail = """<rail>
        <output type="list" description="A list of integers.">
            <integer description="An integer item." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 185μs -> 139μs (33.0% faster)

def test_basic_enum_output():
    """Test enum output schema."""
    rail = """<rail>
        <output type="object">
            <enum name="color" values="red, green, blue" description="Color." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 193μs -> 147μs (30.9% faster)
    props = schema.json_schema["properties"]

def test_basic_messages_and_reask_messages():
    """Test that messages and reask_messages are extracted."""
    rail = """<rail>
        <output type="string" />
        <messages>
            <message role="system">Hello!</message>
            <message role="user">Hi!</message>
        </messages>
        <reask_messages>
            <message role="system">Try again.</message>
        </reask_messages>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 176μs -> 131μs (34.0% faster)

# Edge Test Cases

def test_missing_output_tag_raises():
    """Test that missing <output> raises ValueError."""
    rail = """<rail>
        <messages>
            <message role="system">Hello!</message>
        </messages>
    </rail>"""
    with pytest.raises(ValueError, match="RAIL must contain a output element!"):
        rail_string_to_schema(rail) # 37.0μs -> 35.2μs (5.10% faster)


def test_list_with_multiple_children_raises():
    """Test that <list> with multiple children raises ValueError."""
    rail = """<rail>
        <output type="list">
            <string />
            <integer />
        </output>
    </rail>"""
    with pytest.raises(ValueError, match="<list /> RAIL elements must have precisely 1 child element!"):
        rail_string_to_schema(rail) # 77.7μs -> 75.7μs (2.73% faster)



def test_object_with_no_children_additional_properties_true():
    """Test that object with no children sets additionalProperties to True."""
    rail = """<rail>
        <output type="object" description="Empty object." />
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 158μs -> 128μs (23.1% faster)

def test_enum_with_empty_values():
    """Test enum with empty values attribute."""
    rail = """<rail>
        <output type="object">
            <enum name="empty_enum" values="" description="Empty enum." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 192μs -> 145μs (32.5% faster)

def test_list_with_no_child_items_empty_dict():
    """Test that <list> with no child gives items={}."""
    rail = """<rail>
        <output type="list" />
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 142μs -> 110μs (28.4% faster)

def test_object_with_optional_field():
    """Test object with required and optional fields."""
    rail = """<rail>
        <output type="object">
            <string name="foo" required="true" />
            <integer name="bar" required="false" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 226μs -> 165μs (36.4% faster)

def test_string_with_format():
    """Test <string> with format attribute."""
    rail = """<rail>
        <output type="object">
            <string name="email" format="email" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 186μs -> 143μs (30.2% faster)

def test_date_with_custom_format():
    """Test <date> with custom format and date-format attributes."""
    rail = """<rail>
        <output type="object">
            <date name="dob" format="custom" date-format="%Y-%m-%d" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 196μs -> 152μs (29.5% faster)

# Large Scale Test Cases

def test_large_object_schema():
    """Test large object with many fields."""
    num_fields = 500
    fields = "\n".join(
        [f'<string name="field{i}" description="Field {i}" />' for i in range(num_fields)]
    )
    rail = f"""<rail>
        <output type="object">
            {fields}
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 14.5ms -> 7.85ms (84.8% faster)
    props = schema.json_schema["properties"]
    for i in range(num_fields):
        pass

def test_large_list_schema():
    """Test large list with a complex item schema."""
    rail = """<rail>
        <output type="list">
            <object>
                {}
            </object>
        </output>
    </rail>""".format(
        "\n".join([f'<integer name="num{i}" />' for i in range(100)])
    )
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 3.03ms -> 1.71ms (76.9% faster)
    items = schema.json_schema["items"]
    for i in range(100):
        pass

def test_large_enum_schema():
    """Test enum with many values."""
    values = ",".join([f"v{i}" for i in range(200)])
    rail = f"""<rail>
        <output type="object">
            <enum name="big_enum" values="{values}" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 223μs -> 179μs (24.6% faster)
    enum_values = schema.json_schema["properties"]["big_enum"]["enum"]
    for i in range(200):
        pass


def test_large_nested_object():
    """Test deeply nested object."""
    rail = """<rail>
        <output type="object">
            <object name="level1">
                <object name="level2">
                    <object name="level3">
                        <string name="deep_field" />
                    </object>
                </object>
            </object>
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 284μs -> 210μs (35.1% faster)
    props = schema.json_schema["properties"]
    l1 = props["level1"]["properties"]
    l2 = l1["level2"]["properties"]
    l3 = l2["level3"]["properties"]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import xml.etree.ElementTree as ET

# imports
import pytest
from guardrails.schema.rail_schema import rail_string_to_schema


# Minimal stubs for dependencies used in rail_string_to_schema
class SimpleTypes:
    STRING = "string"
    INTEGER = "integer"
    NUMBER = "number"
    BOOLEAN = "boolean"
    OBJECT = "object"
    ARRAY = "array"

class OutputTypes:
    STRING = "string"
    LIST = "list"
    DICT = "dict"

class GuardExecutionOptions:
    def __init__(self):
        self.messages = None
        self.reask_messages = None

class ProcessedSchema:
    def __init__(self, validators, validator_map, exec_opts):
        self.validators = validators
        self.validator_map = validator_map
        self.exec_opts = exec_opts
        self.json_schema = None
        self.output_type = None
from guardrails.schema.rail_schema import rail_string_to_schema

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

# Basic Test Cases
def test_basic_string_output():
    # Test with a simple string output
    rail = """
    <rail>
      <output type="string" description="A simple string output." />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 142μs -> 112μs (26.4% faster)

def test_basic_object_output():
    # Test with a simple object output with two fields
    rail = """
    <rail>
      <output type="object" description="User info.">
        <string name="username" description="User's name." />
        <integer name="age" description="User's age." />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 228μs -> 168μs (35.7% faster)

def test_basic_list_output():
    # Test with a simple list output of strings
    rail = """
    <rail>
      <output type="list" description="List of tags.">
        <string description="A tag." />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 182μs -> 139μs (30.5% faster)

def test_basic_enum_output():
    # Test with an enum output
    rail = """
    <rail>
      <output type="enum" description="Color">
        <enum values="red, green, blue" description="Color enum." />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 145μs -> 116μs (24.9% faster)
    # The values should be ['red', 'green', 'blue']

def test_basic_messages_and_reask_messages():
    # Test with messages and reask_messages
    rail = """
    <rail>
      <output type="string" />
      <messages>
        <message role="system">Hello!</message>
        <message role="user">What's your name?</message>
      </messages>
      <reask_messages>
        <message role="system">Try again.</message>
      </reask_messages>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 178μs -> 133μs (33.9% faster)

# Edge Test Cases

def test_missing_output_element():
    # Should raise ValueError if output element is missing
    rail = """
    <rail>
      <messages>
        <message role="system">No output!</message>
      </messages>
    </rail>
    """
    with pytest.raises(ValueError, match="RAIL must contain a output element!"):
        rail_string_to_schema(rail) # 36.9μs -> 36.2μs (1.92% faster)

def test_list_with_multiple_children():
    # Should raise ValueError if <list> has more than one child
    rail = """
    <rail>
      <output type="list">
        <string />
        <integer />
      </output>
    </rail>
    """
    with pytest.raises(ValueError, match="must have precisely 1 child"):
        rail_string_to_schema(rail) # 73.3μs -> 71.2μs (2.94% faster)

def test_choice_missing_discriminator():
    # Should raise ValueError if <choice> is missing discriminator
    rail = """
    <rail>
      <output type="choice">
        <case name="A">
          <string name="foo" />
        </case>
      </output>
    </rail>
    """
    with pytest.raises(ValueError, match="must specify a discriminator"):
        rail_string_to_schema(rail) # 73.5μs -> 71.4μs (2.91% faster)

def test_case_missing_name():
    # Should raise ValueError if <case> is missing name
    rail = """
    <rail>
      <output type="choice" discriminator="type">
        <case>
          <string name="foo" />
        </case>
      </output>
    </rail>
    """
    with pytest.raises(ValueError, match="must specify a name"):
        rail_string_to_schema(rail) # 116μs -> 87.5μs (33.1% faster)

def test_object_with_nameless_child():
    # Should skip nameless child and not crash
    rail = """
    <rail>
      <output type="object">
        <string name="foo" />
        <integer />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 702μs -> 653μs (7.58% faster)

def test_object_with_no_children():
    # Should set additionalProperties=True if no children
    rail = """
    <rail>
      <output type="object" />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 153μs -> 121μs (25.7% faster)

def test_list_with_no_children():
    # Should set items to {} if no children
    rail = """
    <rail>
      <output type="list" />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 143μs -> 112μs (27.3% faster)

def test_enum_with_no_values():
    # Should set enum to None or empty
    rail = """
    <rail>
      <output type="enum">
        <enum />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 141μs -> 112μs (25.5% faster)

def test_custom_format_handling():
    # Should handle custom format attributes correctly
    rail = """
    <rail>
      <output type="date" description="A date." format="foo" date-format="%Y-%m-%d" />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 153μs -> 122μs (24.7% faster)

def test_required_false_handling():
    # Should not include field in required if required="false"
    rail = """
    <rail>
      <output type="object">
        <string name="foo" required="false" />
        <string name="bar" required="true" />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 226μs -> 162μs (38.9% faster)

# Large Scale Test Cases

def test_large_object_output():
    # Test with a large object output
    num_fields = 100
    fields = "\n".join([f'<string name="field{i}" />' for i in range(num_fields)])
    rail = f"""
    <rail>
      <output type="object">
        {fields}
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 2.94ms -> 1.62ms (82.1% faster)
    for i in range(num_fields):
        pass

def test_large_list_output():
    # Test with a large list output of objects
    rail = """
    <rail>
      <output type="list">
        <object>
          {}
        </object>
      </output>
    </rail>
    """.format("\n".join([f'<integer name="num{i}" />' for i in range(50)]))
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 1.63ms -> 940μs (73.5% faster)
    for i in range(50):
        pass

def test_large_enum_output():
    # Test with a large enum output
    enum_values = ",".join([f"val{i}" for i in range(200)])
    rail = f"""
    <rail>
      <output type="enum">
        <enum values="{enum_values}" />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 146μs -> 117μs (24.8% faster)
    for i in range(200):
        pass

def test_large_choice_output():
    # Test with a choice output with many cases
    cases = "\n".join([
        f'<case name="case{i}"><string name="field{i}" /></case>'
        for i in range(20)
    ])
    rail = f"""
    <rail>
      <output type="choice" discriminator="type">
        {cases}
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 1.48ms -> 1.08ms (36.6% faster)
    for i in range(20):
        pass

def test_large_messages():
    # Test with many messages
    messages = "\n".join([f'<message role="user">msg{i}</message>' for i in range(50)])
    rail = f"""
    <rail>
      <output type="string" />
      <messages>
        {messages}
      </messages>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 217μs -> 168μs (29.1% faster)
    for i in range(50):
        pass

def test_large_reask_messages():
    # Test with many reask messages
    messages = "\n".join([f'<message role="system">reask{i}</message>' for i in range(50)])
    rail = f"""
    <rail>
      <output type="string" />
      <reask_messages>
        {messages}
      </reask_messages>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 184μs -> 155μs (19.2% faster)
    for i in range(50):
        pass
# 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-rail_string_to_schema-mh1sqc68 and push.