⚡️ Speed up method JsonSchemaTransformer.walk by 730%

.github/workflows/codeflash.yaml

@@ -0,0 +1,37 @@
+name: Codeflash
+
+on:
+  pull_request:
+    paths:
+      # So that this workflow only runs when code within the target module is modified
+      - 'pydantic_ai_slim/pydantic_ai/**'
+  workflow_dispatch:
+
+concurrency:
+  # Any new push to the PR will cancel the previous run, so that only the latest code is optimized
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+
+jobs:
+  optimize:
+    name: Optimize new Python code
+    # Don't run codeflash on codeflash-ai[bot] commits, prevent duplicate optimizations
+    if: ${{ github.actor != 'codeflash-ai[bot]' }}
+    runs-on: ubuntu-latest
+    env:
+      CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }}
+
+    steps:
+      - name: 🛎️ Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: 🐍 Setup UV
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+      - name: 📦 Install Dependencies
+        run: uv sync --all-extras
+      - name: ⚡️Codeflash Optimization
+        run: uv run codeflash

pydantic_ai_slim/pydantic_ai/profiles/_json_schema.py

@@ -45,10 +45,9 @@ def transform(self, schema: JsonSchema) -> JsonSchema:
         return schema
 
     def walk(self) -> JsonSchema:
-        schema = deepcopy(self.schema)
+        schema = {k: v for k, v in self.schema.items() if k != '$defs'}
 
         # First, handle everything but $defs:
-        schema.pop('$defs', None)
         handled = self._handle(schema)
 
         if not self.prefer_inlined_defs and self.defs:
@@ -76,7 +75,10 @@ def _handle(self, schema: JsonSchema) -> JsonSchema:
         nested_refs = 0
         if self.prefer_inlined_defs:
             while ref := schema.get('$ref'):
-                key = re.sub(r'^#/\$defs/', '', ref)
+                if ref.startswith('#/$defs/'):
+                    key = ref[8:]  # len('#/$defs/')
+                else:
+                    key = re.sub(r'^#/\$defs/', '', ref)
                 if key in self.refs_stack:
                     self.recursive_refs.add(key)
                     break  # recursive ref can't be unpacked
@@ -95,8 +97,10 @@ def _handle(self, schema: JsonSchema) -> JsonSchema:
         elif type_ == 'array':
             schema = self._handle_array(schema)
         elif type_ is None:
-            schema = self._handle_union(schema, 'anyOf')
-            schema = self._handle_union(schema, 'oneOf')
+            if 'anyOf' in schema:
+                schema = self._handle_union(schema, 'anyOf')
+            if 'oneOf' in schema:
+                schema = self._handle_union(schema, 'oneOf')
 
         # Apply the base transform
         schema = self.transform(schema)

pyproject.toml

@@ -5,45 +5,6 @@ build-backend = "hatchling.build"
 [tool.hatch.version]
 source = "uv-dynamic-versioning"
 
-[tool.uv-dynamic-versioning]
-vcs = "git"
-style = "pep440"
-bump = true
-
-[project]
-name = "pydantic-ai"
-dynamic = ["version", "dependencies", "optional-dependencies"]
-description = "Agent Framework / shim to use Pydantic with LLMs"
-authors = [
-    { name = "Samuel Colvin", email = "[email protected]" },
-    { name = "Marcelo Trylesinski", email = "[email protected]" },
-    { name = "David Montague", email = "[email protected]" },
-    { name = "Alex Hall", email = "[email protected]" },
-]
-license = "MIT"
-readme = "README.md"
-classifiers = [
-    "Development Status :: 4 - Beta",
-    "Programming Language :: Python",
-    "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
-    "Programming Language :: Python :: 3.12",
-    "Programming Language :: Python :: 3.13",
-    "Intended Audience :: Developers",
-    "Intended Audience :: Information Technology",
-    "License :: OSI Approved :: MIT License",
-    "Operating System :: OS Independent",
-    "Topic :: Internet",
-    "Topic :: Scientific/Engineering :: Artificial Intelligence",
-    "Topic :: Software Development :: Libraries :: Python Modules",
-    "Framework :: Pydantic",
-    "Framework :: Pydantic :: 2",
-]
-requires-python = ">=3.9"
-
 [tool.hatch.metadata.hooks.uv-dynamic-versioning]
 dependencies = [
     "pydantic-ai-slim[openai,vertexai,google,groq,anthropic,mistral,cohere,bedrock,cli,mcp,evals]=={{ version }}",
@@ -54,14 +15,18 @@ examples = ["pydantic-ai-examples=={{ version }}"]
 logfire = ["logfire>=3.11.0"]
 a2a = ["fasta2a>=0.4.1"]
 
-[project.urls]
-Homepage = "https://ai.pydantic.dev"
-Source = "https://github.com/pydantic/pydantic-ai"
-Documentation = "https://ai.pydantic.dev"
-Changelog = "https://github.com/pydantic/pydantic-ai/releases"
 
-[project.scripts]
-pai = "pydantic_ai._cli:cli_exit" # TODO remove this when clai has been out for a while
+[tool.hatch.build.targets.wheel]
+only-include = ["/README.md"]
+
+[tool.hatch.build.targets.sdist]
+include = ["/README.md", "/Makefile", "/tests"]
+
+
+[tool.uv-dynamic-versioning]
+vcs = "git"
+style = "pep440"
+bump = true
 
 [tool.uv.sources]
 pydantic-ai-slim = { workspace = true }
@@ -79,25 +44,6 @@ members = [
     "examples",
 ]
 
-[dependency-groups]
-# dev dependencies are defined in `pydantic-ai-slim/pyproject.toml` to allow for minimal testing
-lint = ["mypy>=1.11.2", "pyright>=1.1.390", "ruff>=0.6.9"]
-docs = [
-    "pydantic-ai[a2a]",
-    "black>=24.10.0",
-    "mkdocs>=1.6.1",
-    "mkdocs-glightbox>=0.4.0",
-    "mkdocs-llmstxt>=0.2.0",
-    "mkdocs-material[imaging]>=9.5.45",
-    "mkdocstrings-python>=1.12.2",
-]
-docs-upload = ["algoliasearch>=4.12.0", "pydantic>=2.10.1"]
-
-[tool.hatch.build.targets.wheel]
-only-include = ["/README.md"]
-
-[tool.hatch.build.targets.sdist]
-include = ["/README.md", "/Makefile", "/tests"]
 
 [tool.ruff]
 line-length = 120
@@ -152,6 +98,7 @@ quote-style = "single"
 "tests/**/*.py" = ["D"]
 "docs/**/*.py" = ["D"]
 
+
 [tool.pyright]
 pythonVersion = "3.12"
 typeCheckingMode = "strict"
@@ -208,6 +155,7 @@ filterwarnings = [
 ]
 
 # https://coverage.readthedocs.io/en/latest/config.html#run
+
 [tool.coverage.run]
 # required to avoid warnings about files created by create_module fixture
 include = [
@@ -250,6 +198,7 @@ exclude_lines = [
     'assert False',
 ]
 
+
 [tool.logfire]
 ignore_no_config = true
 
@@ -260,10 +209,77 @@ format-command = "ruff format --stdin-filename {filename}"
 snap-fix = ["create", "fix"]
 snap = ["create"]
 
+
 [tool.codespell]
 # Ref: https://github.com/codespell-project/codespell#using-a-config-file
 skip = '.git*,*.svg,*.lock,*.css,*.yaml'
 check-hidden = true
 # Ignore "formatting" like **L**anguage
 ignore-regex = '\*\*[A-Z]\*\*[a-z]+\b'
 ignore-words-list = 'asend,aci'
+
+[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "pydantic_ai_slim/pydantic_ai"
+tests-root = "tests"
+test-framework = "pytest"
+ignore-paths = []
+formatter-cmds = ["disabled"]
+
+[project]
+name = "pydantic-ai"
+dynamic = ["version", "dependencies", "optional-dependencies"]
+description = "Agent Framework / shim to use Pydantic with LLMs"
+authors = [
+    { name = "Samuel Colvin", email = "[email protected]" },
+    { name = "Marcelo Trylesinski", email = "[email protected]" },
+    { name = "David Montague", email = "[email protected]" },
+    { name = "Alex Hall", email = "[email protected]" },
+]
+license = "MIT"
+readme = "README.md"
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Information Technology",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Internet",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Framework :: Pydantic",
+    "Framework :: Pydantic :: 2",
+]
+requires-python = ">=3.9"
+
+[project.urls]
+Homepage = "https://ai.pydantic.dev"
+Source = "https://github.com/pydantic/pydantic-ai"
+Documentation = "https://ai.pydantic.dev"
+Changelog = "https://github.com/pydantic/pydantic-ai/releases"
+
+[project.scripts]
+pai = "pydantic_ai._cli:cli_exit" # TODO remove this when clai has been out for a while
+
+[dependency-groups]
+# dev dependencies are defined in `pydantic-ai-slim/pyproject.toml` to allow for minimal testing
+lint = ["mypy>=1.11.2", "pyright>=1.1.390", "ruff>=0.6.9"]
+docs = [
+    "pydantic-ai[a2a]",
+    "black>=24.10.0",
+    "mkdocs>=1.6.1",
+    "mkdocs-glightbox>=0.4.0",
+    "mkdocs-llmstxt>=0.2.0",
+    "mkdocs-material[imaging]>=9.5.45",
+    "mkdocstrings-python>=1.12.2",
+]
+docs-upload = ["algoliasearch>=4.12.0", "pydantic>=2.10.1"]
+