List imports from other example files in their own section

Author: DouweM

docs/evals.md

@@ -60,6 +60,7 @@ from dataclasses import dataclass
 
 from pydantic_evals.evaluators import Evaluator, EvaluatorContext
 from pydantic_evals.evaluators.common import IsInstance
+
 from simple_eval_dataset import dataset
 
 dataset.add_evaluator(IsInstance(type_name='str'))  # (1)!
@@ -618,10 +619,11 @@ You can also write datasets as JSON files:
 ```python {title="generate_dataset_example_json.py" requires="generate_dataset_example.py"}
 from pathlib import Path
 
-from generate_dataset_example import AnswerOutput, MetadataType, QuestionInputs
 from pydantic_evals import Dataset
 from pydantic_evals.generation import generate_dataset
 
+from generate_dataset_example import AnswerOutput, MetadataType, QuestionInputs
+
 
 async def main():
     dataset = await generate_dataset(  # (1)!

docs/output.md

@@ -314,6 +314,7 @@ To use this mode, you can wrap the output type(s) in the [`NativeOutput`][pydant
 
 ```python {title="native_output.py" requires="tool_output.py"}
 from pydantic_ai import Agent, NativeOutput
+
 from tool_output import Fruit, Vehicle
 
 agent = Agent(
@@ -347,6 +348,7 @@ To use this mode, you can wrap the output type(s) in the [`PromptedOutput`][pyda
 from pydantic import BaseModel
 
 from pydantic_ai import Agent, PromptedOutput
+
 from tool_output import Vehicle
 
 

docs/retries.md

@@ -258,6 +258,7 @@ The retry transports work with any provider that accepts a custom HTTP client:
 from pydantic_ai import Agent
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
+
 from smart_retry_example import create_retrying_client
 
 client = create_retrying_client()
@@ -271,6 +272,7 @@ agent = Agent(model)
 from pydantic_ai import Agent
 from pydantic_ai.models.anthropic import AnthropicModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
+
 from smart_retry_example import create_retrying_client
 
 client = create_retrying_client()
@@ -284,6 +286,7 @@ agent = Agent(model)
 from pydantic_ai import Agent
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
+
 from smart_retry_example import create_retrying_client
 
 client = create_retrying_client()
@@ -320,6 +323,7 @@ The retry transports will re-raise the last exception if all retry attempts fail
 from pydantic_ai import Agent
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
+
 from smart_retry_example import create_retrying_client
 
 client = create_retrying_client()

docs/testing.md

@@ -253,6 +253,7 @@ Here's an example of a fixture that overrides the model with `TestModel`:
 import pytest
 
 from pydantic_ai.models.test import TestModel
+
 from weather_app import weather_agent
 
 

docs/toolsets.md

@@ -124,11 +124,12 @@ Toolsets can be composed to dynamically filter which tools are available, modify
 [`CombinedToolset`][pydantic_ai.toolsets.CombinedToolset] takes a list of toolsets and lets them be used as one.
 
 ```python {title="combined_toolset.py" requires="function_toolset.py"}
-from function_toolset import datetime_toolset, weather_toolset
 from pydantic_ai import Agent
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.toolsets import CombinedToolset
 
+from function_toolset import datetime_toolset, weather_toolset
+
 combined_toolset = CombinedToolset([weather_toolset, datetime_toolset])
 
 test_model = TestModel() # (1)!
@@ -149,10 +150,11 @@ _(This example is complete, it can be run "as is")_
 To easily chain different modifications, you can also call [`filtered()`][pydantic_ai.toolsets.AbstractToolset.filtered] on any toolset instead of directly constructing a `FilteredToolset`.
 
 ```python {title="filtered_toolset.py" requires="function_toolset.py,combined_toolset.py"}
-from combined_toolset import combined_toolset
 from pydantic_ai import Agent
 from pydantic_ai.models.test import TestModel
 
+from combined_toolset import combined_toolset
+
 filtered_toolset = combined_toolset.filtered(lambda ctx, tool_def: 'fahrenheit' not in tool_def.name)
 
 test_model = TestModel() # (1)!
@@ -173,11 +175,12 @@ _(This example is complete, it can be run "as is")_
 To easily chain different modifications, you can also call [`prefixed()`][pydantic_ai.toolsets.AbstractToolset.prefixed] on any toolset instead of directly constructing a `PrefixedToolset`.
 
 ```python {title="combined_toolset.py" requires="function_toolset.py"}
-from function_toolset import datetime_toolset, weather_toolset
 from pydantic_ai import Agent
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.toolsets import CombinedToolset
 
+from function_toolset import datetime_toolset, weather_toolset
+
 combined_toolset = CombinedToolset(
     [
         weather_toolset.prefixed('weather'),
@@ -210,10 +213,11 @@ _(This example is complete, it can be run "as is")_
 To easily chain different modifications, you can also call [`renamed()`][pydantic_ai.toolsets.AbstractToolset.renamed] on any toolset instead of directly constructing a `RenamedToolset`.
 
 ```python {title="renamed_toolset.py" requires="function_toolset.py,combined_toolset.py"}
-from combined_toolset import combined_toolset
 from pydantic_ai import Agent
 from pydantic_ai.models.test import TestModel
 
+from combined_toolset import combined_toolset
+
 renamed_toolset = combined_toolset.renamed(
     {
         'current_time': 'datetime_now',
@@ -250,6 +254,7 @@ from dataclasses import replace
 
 from pydantic_ai import Agent, RunContext, ToolDefinition
 from pydantic_ai.models.test import TestModel
+
 from renamed_toolset import renamed_toolset
 
 descriptions = {
@@ -330,10 +335,11 @@ To easily chain different modifications, you can also call [`approval_required()
 See the [Human-in-the-Loop Tool Approval](tools.md#human-in-the-loop-tool-approval) documentation for more information on how to handle agent runs that call tools that require approval and how to pass in the results.
 
 ```python {title="approval_required_toolset.py" requires="function_toolset.py,combined_toolset.py,renamed_toolset.py,prepared_toolset.py"}
-from prepared_toolset import prepared_toolset
 from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults
 from pydantic_ai.models.test import TestModel
 
+from prepared_toolset import prepared_toolset
+
 approval_required_toolset = prepared_toolset.approval_required(lambda ctx, tool_def, tool_args: tool_def.name.startswith('temperature'))
 
 test_model = TestModel(call_tools=['temperature_celsius', 'temperature_fahrenheit']) # (1)!
@@ -391,11 +397,12 @@ import asyncio
 
 from typing_extensions import Any
 
-from prepared_toolset import prepared_toolset
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.toolsets import ToolsetTool, WrapperToolset
 
+from prepared_toolset import prepared_toolset
+
 LOG = []
 
 class LoggingToolset(WrapperToolset):
@@ -483,11 +490,12 @@ print(repr(result.output))
 Next, let's define a function that represents a hypothetical "run agent" API endpoint that can be called by the frontend and takes a list of messages to send to the model, a list of frontend tool definitions, and optional deferred tool results. This is where `ExternalToolset`, `DeferredToolRequests`, and `DeferredToolResults` come in:
 
 ```python {title="deferred_toolset_api.py" requires="deferred_toolset_agent.py"}
-from deferred_toolset_agent import PersonalizedGreeting, agent
 from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDefinition
 from pydantic_ai.messages import ModelMessage
 from pydantic_ai.toolsets import ExternalToolset
 
+from deferred_toolset_agent import PersonalizedGreeting, agent
+
 
 def run_agent(
     messages: list[ModelMessage] = [],
@@ -511,7 +519,6 @@ def run_agent(
 Now, imagine that the code below is implemented on the frontend, and `run_agent` stands in for an API call to the backend that runs the agent. This is where we actually execute the deferred tool calls and start a new run with the new result included:
 
 ```python {title="deferred_tools.py" requires="deferred_toolset_agent.py,deferred_toolset_api.py"}
-from deferred_toolset_api import run_agent
 from pydantic_ai import (
     DeferredToolRequests,
     DeferredToolResults,
@@ -520,6 +527,8 @@ from pydantic_ai import (
 )
 from pydantic_ai.messages import ModelMessage, ModelRequest, UserPromptPart
 
+from deferred_toolset_api import run_agent
+
 frontend_tool_definitions = [
     ToolDefinition(
         name='get_preferred_language',
@@ -592,10 +601,11 @@ By default, the function will be called again ahead of each agent run step. If y
 from dataclasses import dataclass
 from typing import Literal
 
-from function_toolset import datetime_toolset, weather_toolset
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.models.test import TestModel
 
+from function_toolset import datetime_toolset, weather_toolset
+
 
 @dataclass
 class ToggleableDeps:

pydantic_ai_slim/pydantic_ai/output.py

@@ -133,6 +133,7 @@ class NativeOutput(Generic[OutputDataT]):
     Example:
     ```python {title="native_output.py" requires="tool_output.py"}
     from pydantic_ai import Agent, NativeOutput
+
     from tool_output import Fruit, Vehicle
 
     agent = Agent(
@@ -181,6 +182,7 @@ class PromptedOutput(Generic[OutputDataT]):
     from pydantic import BaseModel
 
     from pydantic_ai import Agent, PromptedOutput
+
     from tool_output import Vehicle
 
 

tests/test_examples.py

@@ -7,7 +7,7 @@
 import ssl
 import sys
 from collections.abc import AsyncIterator, Iterable, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from inspect import FrameInfo
 from io import StringIO
 from pathlib import Path
@@ -57,14 +57,18 @@
 code_examples: dict[str, CodeExample] = {}
 
 
+@dataclass
 class ExamplesConfig(BaseExamplesConfig):
+    known_first_party: list[str] = field(default_factory=list)
+    known_local_folder: list[str] = field(default_factory=list)
+
     def ruff_config(self) -> tuple[str, ...]:
         config = super().ruff_config()
-        return (
-            *config,
-            '--config',
-            'lint.isort.known-first-party = ["pydantic_ai", "pydantic_evals", "pydantic_graph"]',
-        )
+        if self.known_first_party:
+            config = (*config, '--config', f'lint.isort.known-first-party = {self.known_first_party}')
+        if self.known_local_folder:
+            config = (*config, '--config', f'lint.isort.known-local-folder = {self.known_local_folder}')
+        return config
 
 
 def find_filter_examples() -> Iterable[ParameterSet]:
@@ -179,8 +183,10 @@ def print(self, *args: Any, **kwargs: Any) -> None:
     if opt_test.startswith('skip') and opt_lint.startswith('skip'):
         pytest.skip('both running code and lint skipped')
 
+    known_local_folder: list[str] = []
     if requires:
         for req in requires.split(','):
+            known_local_folder.append(Path(req).stem)
             if ex := code_examples.get(req):
                 (tmp_path_cwd / req).write_text(ex.source)
             else:  # pragma: no cover
@@ -205,6 +211,8 @@ def print(self, *args: Any, **kwargs: Any) -> None:
         isort=True,
         upgrade=True,
         quotes='single',
+        known_first_party=['pydantic_ai', 'pydantic_evals', 'pydantic_graph'],
+        known_local_folder=known_local_folder,
     )
     eval_example.print_callback = print_callback
     eval_example.include_print = custom_include_print