lint fixes after merge

Author: pirate

src/stagehand/_pydantic_extract.py

@@ -4,7 +4,8 @@
 
 import inspect
 import logging
-from typing import Any, Dict, Type
+from typing import Any, Dict, Type, cast
+from typing_extensions import Literal
 
 from pydantic import BaseModel, ConfigDict
 
@@ -58,14 +59,17 @@ def validate_extract_response(
 
 @lru_cache(maxsize=None)
 def _validation_schema(schema: Type[BaseModel], strict_response_validation: bool) -> Type[BaseModel]:
-    extra_behavior = "forbid" if strict_response_validation else "allow"
-    validation_schema = type(
-        f"{schema.__name__}ExtractValidation",
-        (schema,),
-        {
-            "__module__": schema.__module__,
-            "model_config": ConfigDict(extra=extra_behavior),
-        },
+    extra_behavior: Literal["allow", "forbid"] = "forbid" if strict_response_validation else "allow"
+    validation_schema = cast(
+        Type[BaseModel],
+        type(
+            f"{schema.__name__}ExtractValidation",
+            (schema,),
+            {
+                "__module__": schema.__module__,
+                "model_config": ConfigDict(extra=extra_behavior),
+            },
+        ),
     )
     validation_schema.model_rebuild(force=True)
     return validation_schema

src/stagehand/_session_extract.py

@@ -2,15 +2,15 @@
 
 from __future__ import annotations
 
-from typing import Any, cast
+from typing import Any, Mapping, cast
+from typing_extensions import Unpack
 
 import httpx
-from typing_extensions import Unpack
 
-from ._pydantic_extract import is_pydantic_model, pydantic_model_to_json_schema, validate_extract_response
-from .session import AsyncSession, Session
-from ._types import Body, Headers, NotGiven, Query, not_given
 from .types import session_extract_params
+from ._types import Body, Query, Headers, NotGiven, not_given
+from .session import Session, AsyncSession
+from ._pydantic_extract import is_pydantic_model, validate_extract_response, pydantic_model_to_json_schema
 from .types.session_extract_response import SessionExtractResponse
 
 _ORIGINAL_SESSION_EXTRACT = Session.extract
@@ -104,23 +104,22 @@ async def _async_extract(  # type: ignore[override, misc]
 
 
 def _with_schema(
-    params: session_extract_params.SessionExtractParamsNonStreaming,
+    params: Mapping[str, object],
     schema: dict[str, object] | type | None,
 ) -> session_extract_params.SessionExtractParamsNonStreaming:
     api_params = dict(params)
     if schema is not None:
         api_params["schema"] = cast(Any, schema)
     return cast(session_extract_params.SessionExtractParamsNonStreaming, api_params)
 
-
 _sync_extract.__module__ = _ORIGINAL_SESSION_EXTRACT.__module__
 _sync_extract.__name__ = _ORIGINAL_SESSION_EXTRACT.__name__
 _sync_extract.__qualname__ = _ORIGINAL_SESSION_EXTRACT.__qualname__
 _sync_extract.__doc__ = _ORIGINAL_SESSION_EXTRACT.__doc__
-_sync_extract.__stagehand_pydantic_extract_patch__ = True
+setattr(_sync_extract, "__stagehand_pydantic_extract_patch__", True)  # noqa: B010
 
 _async_extract.__module__ = _ORIGINAL_ASYNC_SESSION_EXTRACT.__module__
 _async_extract.__name__ = _ORIGINAL_ASYNC_SESSION_EXTRACT.__name__
 _async_extract.__qualname__ = _ORIGINAL_ASYNC_SESSION_EXTRACT.__qualname__
 _async_extract.__doc__ = _ORIGINAL_ASYNC_SESSION_EXTRACT.__doc__
-_async_extract.__stagehand_pydantic_extract_patch__ = True
+setattr(_async_extract, "__stagehand_pydantic_extract_patch__", True)  # noqa: B010

src/stagehand/resources/sessions_helpers.py

@@ -6,26 +6,26 @@
 
 import httpx
 
-from .._session_extract import install_pydantic_extract_patch
-from .._compat import cached_property
-from .._response import (
-    async_to_raw_response_wrapper,
-    async_to_streamed_response_wrapper,
-    to_raw_response_wrapper,
-    to_streamed_response_wrapper,
-)
+from ..types import session_start_params
 from .._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
+from .._compat import cached_property
 from ..session import Session, AsyncSession
-from ..types import session_start_params
-from ..types.session_start_response import SessionStartResponse
 from .sessions import (
-    AsyncSessionsResource,
-    AsyncSessionsResourceWithRawResponse,
-    AsyncSessionsResourceWithStreamingResponse,
     SessionsResource,
+    AsyncSessionsResource,
     SessionsResourceWithRawResponse,
+    AsyncSessionsResourceWithRawResponse,
     SessionsResourceWithStreamingResponse,
+    AsyncSessionsResourceWithStreamingResponse,
 )
+from .._response import (
+    to_raw_response_wrapper,
+    to_streamed_response_wrapper,
+    async_to_raw_response_wrapper,
+    async_to_streamed_response_wrapper,
+)
+from .._session_extract import install_pydantic_extract_patch
+from ..types.session_start_response import SessionStartResponse
 
 install_pydantic_extract_patch()
 

tests/test_session_extract_pydantic.py

@@ -4,15 +4,15 @@
 
 import os
 import json
-from typing import cast
+from typing import Any, cast
 
 import httpx
 import pytest
-from pydantic import BaseModel
 from respx import MockRouter
+from pydantic import BaseModel
 from respx.models import Call
 
-from stagehand import AsyncStagehand, Stagehand
+from stagehand import Stagehand, AsyncStagehand
 
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
 
@@ -45,7 +45,7 @@ def test_session_extract_accepts_pydantic_schema(respx_mock: MockRouter, client:
     )
 
     session = client.sessions.start(model_name="openai/gpt-5-nano")
-    response = session.extract(instruction="extract the user", schema=ExtractedUser)
+    response = session.extract(instruction="extract the user", schema=cast(Any, ExtractedUser))
 
     assert isinstance(response.data.result, ExtractedUser)
     assert response.data.result.user_name == "Ada"
@@ -79,7 +79,7 @@ def test_session_extract_allows_extra_fields_when_client_is_non_strict(
     )
 
     session = client.sessions.start(model_name="openai/gpt-5-nano")
-    response = session.extract(instruction="extract the user", schema=ExtractedName)
+    response = session.extract(instruction="extract the user", schema=cast(Any, ExtractedName))
 
     assert isinstance(response.data.result, ExtractedName)
     assert response.data.result.user_name == "Ada"
@@ -107,7 +107,7 @@ def test_session_extract_rejects_extra_fields_when_client_is_strict(
     )
 
     session = client.sessions.start(model_name="openai/gpt-5-nano")
-    response = session.extract(instruction="extract the user", schema=ExtractedName)
+    response = session.extract(instruction="extract the user", schema=cast(Any, ExtractedName))
 
     assert response.data.result == {"userName": "Ada", "favoriteColor": "blue"}
 
@@ -133,7 +133,7 @@ async def test_async_session_extract_accepts_pydantic_schema(
     )
 
     session = await async_client.sessions.start(model_name="openai/gpt-5-nano")
-    response = await session.extract(instruction="extract the user", schema=ExtractedUser)
+    response = await session.extract(instruction="extract the user", schema=cast(Any, ExtractedUser))
 
     assert isinstance(response.data.result, ExtractedUser)
     assert response.data.result.user_name == "Grace"
@@ -167,7 +167,7 @@ async def test_async_session_extract_allows_extra_fields_when_client_is_non_stri
     )
 
     session = await async_client.sessions.start(model_name="openai/gpt-5-nano")
-    response = await session.extract(instruction="extract the user", schema=ExtractedName)
+    response = await session.extract(instruction="extract the user", schema=cast(Any, ExtractedName))
 
     assert isinstance(response.data.result, ExtractedName)
     assert response.data.result.user_name == "Grace"
@@ -195,6 +195,6 @@ async def test_async_session_extract_rejects_extra_fields_when_client_is_strict(
     )
 
     session = await async_client.sessions.start(model_name="openai/gpt-5-nano")
-    response = await session.extract(instruction="extract the user", schema=ExtractedName)
+    response = await session.extract(instruction="extract the user", schema=cast(Any, ExtractedName))
 
     assert response.data.result == {"userName": "Grace", "favoriteColor": "green"}