⚡️ Speed up function `flatten_schema` by 11% in PR #7741 (`mcp-server-backend`)

[codeflash-ai[bot]]

⚡️ This pull request contains optimizations for PR #7741

If you approve this dependent PR, these changes will be merged into the original PR branch mcp-server-backend.

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

---

📄 11% (0.11x) speedup for flatten_schema in src/backend/base/langflow/io/schema.py

⏱️ Runtime : 803 microseconds → 724 microseconds (best of 235 runs)

📝 Explanation and details

To optimize the given Python program for faster runtime, I'll make a few enhancements to reduce time complexity. The primary areas of improvement focus on avoiding repeated operations and optimizing the recursion. Here's the improved version of the program.

Summary of Changes.

1. _resolve_if_ref function simplified: replaced the loop with a single reference resolution check inside _walk to avoid the need for looping twice over $ref.
2. Recursion Execution: Reduced deep nesting and ensured schema is resolved earlier to prevent multiple lookups.
3. Optimized Set Operations: Replaced creating set inside recursion with a set created outside "_walk".
4. Inlined-Leaf Dictionary Construction: The comprehension inside _walk for creating the leaf object is now written as a single statement, making it concise and slightly faster.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from typing import Any

# imports
import pytest  # used for our unit tests
from langflow.io.schema import flatten_schema

# unit tests

def test_flat_schema():
    # Basic flat schema with no nested objects or arrays
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "number"},
            "isActive": {"type": "boolean"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_schema_with_defs_no_ref():
    # Schema with $defs but no $ref
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "number"}
        },
        "$defs": {
            "address": {
                "type": "object",
                "properties": {
                    "street": {"type": "string"},
                    "city": {"type": "string"}
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "number"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_schema_with_ref():
    # Schema with properties that reference definitions in $defs
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "homeAddress": {"$ref": "#/$defs/address"}
        },
        "$defs": {
            "address": {
                "type": "object",
                "properties": {
                    "street": {"type": "string"},
                    "city": {"type": "string"}
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "homeAddress.street": {"type": "string"},
            "homeAddress.city": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_nested_objects():
    # Schema with nested objects within properties
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "address": {
                "type": "object",
                "properties": {
                    "street": {"type": "string"},
                    "city": {"type": "string"}
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "address.street": {"type": "string"},
            "address.city": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_arrays_of_objects():
    # Schema with properties that are arrays of objects
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "addresses": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "street": {"type": "string"},
                        "city": {"type": "string"}
                    }
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "addresses[0].street": {"type": "string"},
            "addresses[0].city": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_required_properties():
    # Schema with required properties
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "number"}
        },
        "required": ["name"]
    }
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "number"}
        },
        "required": ["name"]
    }
    codeflash_output = flatten_schema(schema)

def test_complex_nested_structures():
    # Complex nested structures with mixed types
    schema = {
        "type": "object",
        "properties": {
            "user": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "contacts": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "type": {"type": "string"},
                                "details": {"$ref": "#/$defs/contactDetails"}
                            }
                        }
                    }
                }
            }
        },
        "$defs": {
            "contactDetails": {
                "type": "object",
                "properties": {
                    "email": {"type": "string"},
                    "phone": {"type": "string"}
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "user.name": {"type": "string"},
            "user.contacts[0].type": {"type": "string"},
            "user.contacts[0].details.email": {"type": "string"},
            "user.contacts[0].details.phone": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_empty_schema():
    # Edge case: Empty schema
    schema = {}
    expected = {}
    codeflash_output = flatten_schema(schema)

def test_empty_properties():
    # Edge case: Schema with an empty properties object
    schema = {
        "type": "object",
        "properties": {}
    }
    expected = {
        "type": "object",
        "properties": {}
    }
    codeflash_output = flatten_schema(schema)

def test_large_scale_schema():
    # Large scale schema with multiple nested levels and numerous properties
    schema = {
        "type": "object",
        "properties": {
            "level1": {
                "type": "object",
                "properties": {
                    "level2": {
                        "type": "object",
                        "properties": {
                            "level3": {
                                "type": "object",
                                "properties": {
                                    "level4": {
                                        "type": "object",
                                        "properties": {
                                            "name": {"type": "string"},
                                            "details": {
                                                "type": "object",
                                                "properties": {
                                                    "age": {"type": "number"},
                                                    "address": {
                                                        "type": "object",
                                                        "properties": {
                                                            "street": {"type": "string"},
                                                            "city": {"type": "string"}
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "level1.level2.level3.level4.name": {"type": "string"},
            "level1.level2.level3.level4.details.age": {"type": "number"},
            "level1.level2.level3.level4.details.address.street": {"type": "string"},
            "level1.level2.level3.level4.details.address.city": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_large_number_of_properties():
    # Large scale schema with a large number of properties and definitions
    schema = {
        "type": "object",
        "properties": {
            **{f"prop{i}": {"type": "string"} for i in range(1, 1001)}
        },
        "$defs": {
            **{f"def{i}": {"type": "object", "properties": {f"attr{i}": {"type": "string"}}} for i in range(1, 101)}
        }
    }
    expected = {
        "type": "object",
        "properties": {
            **{f"prop{i}": {"type": "string"} for i in range(1, 1001)}
        }
    }
    codeflash_output = flatten_schema(schema)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import Any

# imports
import pytest  # used for our unit tests
from langflow.io.schema import flatten_schema

# unit tests

def test_flatten_schema_basic_flat():
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_with_ref():
    schema = {
        "$defs": {
            "address": {
                "type": "object",
                "properties": {
                    "street": {"type": "string"},
                    "city": {"type": "string"}
                }
            }
        },
        "type": "object",
        "properties": {
            "home_address": {"$ref": "#/$defs/address"}
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "home_address.street": {"type": "string"},
            "home_address.city": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_nested_objects():
    schema = {
        "type": "object",
        "properties": {
            "person": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "age": {"type": "integer"}
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "person.name": {"type": "string"},
            "person.age": {"type": "integer"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_with_arrays():
    schema = {
        "type": "object",
        "properties": {
            "tags": {
                "type": "array",
                "items": {"type": "string"}
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "tags[0]": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_mixed_nested():
    schema = {
        "type": "object",
        "properties": {
            "department": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "employees": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "name": {"type": "string"},
                                "age": {"type": "integer"}
                            }
                        }
                    }
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "department.name": {"type": "string"},
            "department.employees[0].name": {"type": "string"},
            "department.employees[0].age": {"type": "integer"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_required_fields():
    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer"}
        },
        "required": ["name"]
    }
    expected = {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer"}
        },
        "required": ["name"]
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_edge_empty():
    schema = {}
    expected = {}
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_edge_unknown_type():
    schema = {
        "type": "object",
        "properties": {
            "unknown": {"type": "unknown"}
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "unknown": {"type": "unknown"}
        }
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_large_scale():
    schema = {
        "type": "object",
        "properties": {f"prop{i}": {"type": "string"} for i in range(1000)}
    }
    expected = {
        "type": "object",
        "properties": {f"prop{i}": {"type": "string"} for i in range(1000)}
    }
    codeflash_output = flatten_schema(schema)

def test_flatten_schema_deeply_nested():
    schema = {
        "type": "object",
        "properties": {
            "level1": {
                "type": "object",
                "properties": {
                    "level2": {
                        "type": "object",
                        "properties": {
                            "level3": {
                                "type": "object",
                                "properties": {
                                    "finalProp": {"type": "string"}
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    expected = {
        "type": "object",
        "properties": {
            "level1.level2.level3.finalProp": {"type": "string"}
        }
    }
    codeflash_output = flatten_schema(schema)
# 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-pr7741-2025-04-24T16.13.13 and push.