Update tests, pass CI

Author: tode-rl

scripts/generate_examples_md.py

@@ -11,10 +11,10 @@
 import re
 import sys
 import argparse
-from typing import Any
+from typing import Any, cast
 from pathlib import Path
 
-import frontmatter  # type: ignore[import-untyped]
+import frontmatter  # type: ignore[import-not-found, import-untyped]
 
 ROOT = Path(__file__).parent.parent
 EXAMPLES_DIR = ROOT / "examples"
@@ -38,7 +38,7 @@ def parse_example(path: Path) -> dict[str, Any]:
         raise ValueError(f"{path}: docstring must start with frontmatter (---)")
 
     try:
-        post = frontmatter.loads(docstring)
+        post = cast(Any, frontmatter).loads(docstring)
         return dict(post.metadata)
     except Exception as e:
         raise ValueError(f"{path}: invalid frontmatter: {e}") from e

src/runloop_api_client/_utils/_compat.py

@@ -2,7 +2,7 @@
 
 import sys
 import typing_extensions
-from typing import Any, Type, Union, Literal, Optional
+from typing import Any, Type, Union, Literal, Optional, cast
 from datetime import date, datetime
 from typing_extensions import get_args as _get_args, get_origin as _get_origin
 
@@ -34,7 +34,8 @@ def is_typeddict(tp: Type[Any]) -> bool:
 
 
 def is_literal_type(tp: Type[Any]) -> bool:
-    return get_origin(tp) in _LITERAL_TYPES
+    origin = get_origin(tp)
+    return cast(Any, origin) in _LITERAL_TYPES
 
 
 def parse_date(value: Union[date, StrBytesIntFloat]) -> date:

src/runloop_api_client/_utils/_utils.py

@@ -372,8 +372,7 @@ def file_from_path(path: str) -> FileTypes:
 def get_required_header(headers: HeadersLike, header: str) -> str:
     lower_header = header.lower()
     if is_mapping_t(headers):
-        # mypy doesn't understand the type narrowing here
-        for k, v in headers.items():  # type: ignore
+        for k, v in cast(Mapping[str, object], headers).items():
             if k.lower() == lower_header and isinstance(v, str):
                 return v
 

src/runloop_api_client/lib/polling_async.py

@@ -1,6 +1,7 @@
 import time
 import asyncio
 from typing import Union, TypeVar, Callable, Optional, Awaitable
+from contextlib import suppress
 
 from .polling import PollingConfig, PollingTimeout
 from .cancellation import CancellationToken
@@ -67,13 +68,13 @@ async def async_poll_until(
 
         # Cancellable async sleep
         if cancellation_token is not None:
+            wait_task = asyncio.create_task(cancellation_token.async_event.wait())
             try:
-                await asyncio.wait_for(
-                    cancellation_token.async_event.wait(),
-                    timeout=config.interval_seconds,
-                )
+                await asyncio.wait_for(wait_task, timeout=config.interval_seconds)
                 cancellation_token.raise_if_cancelled()
             except asyncio.TimeoutError:
-                pass  # Normal sleep completion
+                wait_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await wait_task
         else:
             await asyncio.sleep(config.interval_seconds)

src/runloop_api_client/resources/axons/axons.py

@@ -16,7 +16,6 @@
     AsyncSqlResourceWithStreamingResponse,
 )
 from ...types import axon_list_params, axon_create_params, axon_publish_params
-from ...types.axons import axon_subscribe_sse_params
 from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
 from ..._utils import path_template, maybe_transform, async_maybe_transform
 from ..._compat import cached_property
@@ -30,6 +29,7 @@
 from ..._constants import RAW_RESPONSE_HEADER
 from ..._streaming import Stream, AsyncStream, ReconnectingStream, AsyncReconnectingStream
 from ...pagination import SyncAxonsCursorIDPage, AsyncAxonsCursorIDPage
+from ...types.axons import axon_subscribe_sse_params
 from ..._base_client import AsyncPaginator, make_request_options
 from ...types.axon_view import AxonView
 from ...types.axon_event_view import AxonEventView

src/runloop_api_client/resources/blueprints.py

@@ -28,8 +28,8 @@
 from .._exceptions import RunloopError
 from ..lib.polling import PollingConfig, poll_until
 from .._base_client import AsyncPaginator, make_request_options
-from ..lib.polling_async import async_poll_until
 from ..lib.cancellation import CancellationToken
+from ..lib.polling_async import async_poll_until
 from .._utils._validation import ValidationNotification
 from ..types.blueprint_view import BlueprintView
 from ..types.blueprint_preview_view import BlueprintPreviewView

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -97,8 +97,8 @@
     DiskSnapshotsResourceWithStreamingResponse,
     AsyncDiskSnapshotsResourceWithStreamingResponse,
 )
-from ...lib.polling_async import async_poll_until
 from ...lib.cancellation import CancellationToken
+from ...lib.polling_async import async_poll_until
 from ...types.devbox_view import DevboxView
 from ...types.tunnel_view import TunnelView
 from ...types.shared_params.mount import Mount

src/runloop_api_client/resources/devboxes/disk_snapshots.py

@@ -21,8 +21,8 @@
 from ...lib.polling import PollingConfig, poll_until
 from ..._base_client import AsyncPaginator, make_request_options
 from ...types.devboxes import disk_snapshot_list_params, disk_snapshot_update_params
-from ...lib.polling_async import async_poll_until
 from ...lib.cancellation import CancellationToken
+from ...lib.polling_async import async_poll_until
 from ...types.devbox_snapshot_view import DevboxSnapshotView
 from ...types.devboxes.devbox_snapshot_async_status_view import DevboxSnapshotAsyncStatusView
 

src/runloop_api_client/resources/devboxes/executions.py

@@ -32,8 +32,8 @@
     execution_stream_stderr_updates_params,
     execution_stream_stdout_updates_params,
 )
-from ...lib.polling_async import async_poll_until
 from ...lib.cancellation import CancellationToken
+from ...lib.polling_async import async_poll_until
 from ...types.devbox_send_std_in_result import DevboxSendStdInResult
 from ...types.devbox_execution_detail_view import DevboxExecutionDetailView
 from ...types.devboxes.execution_update_chunk import ExecutionUpdateChunk

src/runloop_api_client/resources/scenarios/runs.py

@@ -27,8 +27,8 @@
 from ...lib.polling import PollingConfig, poll_until
 from ..._base_client import AsyncPaginator, make_request_options
 from ...types.scenarios import run_list_params
-from ...lib.polling_async import async_poll_until
 from ...lib.cancellation import CancellationToken
+from ...lib.polling_async import async_poll_until
 from ...types.scenario_run_view import ScenarioRunView
 
 __all__ = ["RunsResource", "AsyncRunsResource"]
@@ -429,6 +429,7 @@ def score_and_complete(
         id: str,
         *,
         polling_config: PollingConfig | None = None,
+        cancellation_token: CancellationToken | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -441,6 +442,7 @@ def score_and_complete(
         Args:
             id: The ID of the scenario run to score, wait for, and complete
             polling_config: Optional polling configuration
+            cancellation_token: Token to cancel the wait operation
             extra_headers: Send extra headers
             extra_query: Add additional query parameters to the request
             extra_body: Add additional JSON properties to the request
@@ -456,6 +458,7 @@ def score_and_complete(
         self.score_and_await(
             id,
             polling_config=polling_config,
+            cancellation_token=cancellation_token,
             extra_headers=extra_headers,
             extra_query=extra_query,
             extra_body=extra_body,
@@ -866,6 +869,7 @@ async def score_and_complete(
         id: str,
         *,
         polling_config: PollingConfig | None = None,
+        cancellation_token: CancellationToken | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -878,6 +882,7 @@ async def score_and_complete(
         Args:
             id: The ID of the scenario run to score, wait for, and complete
             polling_config: Optional polling configuration
+            cancellation_token: Token to cancel the wait operation
             extra_headers: Send extra headers
             extra_query: Add additional query parameters to the request
             extra_body: Add additional JSON properties to the request
@@ -893,6 +898,7 @@ async def score_and_complete(
         await self.score_and_await(
             id,
             polling_config=polling_config,
+            cancellation_token=cancellation_token,
             extra_headers=extra_headers,
             extra_query=extra_query,
             extra_body=extra_body,