fix(mypy): resolve boto3 S3 typing issues in files.py

Add boto3-stubs[s3] dev dependency and use TYPE_CHECKING with cast() to
provide proper S3Client type hints while avoiding boto3-stubs' overly
strict overload definitions.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Ashwin Bharambe 2025-10-27 23:35:31 -07:00
parent 3c81b23fbe
commit 9c877730ca
3 changed files with 30 additions and 13 deletions

View file

@ -67,8 +67,8 @@ dev = [
"pytest-cov", "pytest-cov",
"pytest-html", "pytest-html",
"pytest-json-report", "pytest-json-report",
"pytest-socket", # For blocking network access in unit tests "pytest-socket", # For blocking network access in unit tests
"nbval", # For notebook testing "nbval", # For notebook testing
"black", "black",
"ruff", "ruff",
"mypy", "mypy",
@ -78,9 +78,9 @@ dev = [
"pandas-stubs", "pandas-stubs",
"types-psutil", "types-psutil",
"types-tqdm", "types-tqdm",
"boto3-stubs", "boto3-stubs[s3]",
"pre-commit", "pre-commit",
"ruamel.yaml", # needed for openapi generator "ruamel.yaml", # needed for openapi generator
] ]
# These are the dependencies required for running unit tests. # These are the dependencies required for running unit tests.
unit = [ unit = [

View file

@ -6,12 +6,15 @@
import uuid import uuid
from datetime import UTC, datetime from datetime import UTC, datetime
from typing import Annotated, Any from typing import TYPE_CHECKING, Annotated, Any, cast
import boto3 import boto3
from botocore.exceptions import BotoCoreError, ClientError, NoCredentialsError from botocore.exceptions import BotoCoreError, ClientError, NoCredentialsError
from fastapi import Depends, File, Form, Response, UploadFile from fastapi import Depends, File, Form, Response, UploadFile
if TYPE_CHECKING:
from mypy_boto3_s3.client import S3Client
from llama_stack.apis.common.errors import ResourceNotFoundError from llama_stack.apis.common.errors import ResourceNotFoundError
from llama_stack.apis.common.responses import Order from llama_stack.apis.common.responses import Order
from llama_stack.apis.files import ( from llama_stack.apis.files import (
@ -34,7 +37,7 @@ from .config import S3FilesImplConfig
# TODO: provider data for S3 credentials # TODO: provider data for S3 credentials
def _create_s3_client(config: S3FilesImplConfig) -> boto3.client: def _create_s3_client(config: S3FilesImplConfig) -> S3Client:
try: try:
s3_config = { s3_config = {
"region_name": config.region, "region_name": config.region,
@ -52,13 +55,13 @@ def _create_s3_client(config: S3FilesImplConfig) -> boto3.client:
} }
) )
return boto3.client("s3", **s3_config) return cast("S3Client", boto3.client("s3", **s3_config)) # type: ignore[call-overload]
except (BotoCoreError, NoCredentialsError) as e: except (BotoCoreError, NoCredentialsError) as e:
raise RuntimeError(f"Failed to initialize S3 client: {e}") from e raise RuntimeError(f"Failed to initialize S3 client: {e}") from e
async def _create_bucket_if_not_exists(client: boto3.client, config: S3FilesImplConfig) -> None: async def _create_bucket_if_not_exists(client: S3Client, config: S3FilesImplConfig) -> None:
try: try:
client.head_bucket(Bucket=config.bucket_name) client.head_bucket(Bucket=config.bucket_name)
except ClientError as e: except ClientError as e:
@ -76,7 +79,7 @@ async def _create_bucket_if_not_exists(client: boto3.client, config: S3FilesImpl
else: else:
client.create_bucket( client.create_bucket(
Bucket=config.bucket_name, Bucket=config.bucket_name,
CreateBucketConfiguration={"LocationConstraint": config.region}, CreateBucketConfiguration=cast(Any, {"LocationConstraint": config.region}),
) )
except ClientError as create_error: except ClientError as create_error:
raise RuntimeError( raise RuntimeError(
@ -128,7 +131,7 @@ class S3FilesImpl(Files):
def __init__(self, config: S3FilesImplConfig, policy: list[AccessRule]) -> None: def __init__(self, config: S3FilesImplConfig, policy: list[AccessRule]) -> None:
self._config = config self._config = config
self.policy = policy self.policy = policy
self._client: boto3.client | None = None self._client: S3Client | None = None
self._sql_store: AuthorizedSqlStore | None = None self._sql_store: AuthorizedSqlStore | None = None
def _now(self) -> int: def _now(self) -> int:
@ -184,7 +187,7 @@ class S3FilesImpl(Files):
pass pass
@property @property
def client(self) -> boto3.client: def client(self) -> S3Client:
assert self._client is not None, "Provider not initialized" assert self._client is not None, "Provider not initialized"
return self._client return self._client

18
uv.lock generated
View file

@ -410,6 +410,11 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/83/8a/d14e63701c4e869f1d37ba5657f9821961616b98a30074f20b559c071fb6/boto3_stubs-1.40.60-py3-none-any.whl", hash = "sha256:1ea7f9dbabc7f9ac8477646c12cc51ef49af6b24d53cc2ae8cf6fa6bed6a995a", size = 69746, upload-time = "2025-10-27T19:49:05.619Z" }, { url = "https://files.pythonhosted.org/packages/83/8a/d14e63701c4e869f1d37ba5657f9821961616b98a30074f20b559c071fb6/boto3_stubs-1.40.60-py3-none-any.whl", hash = "sha256:1ea7f9dbabc7f9ac8477646c12cc51ef49af6b24d53cc2ae8cf6fa6bed6a995a", size = 69746, upload-time = "2025-10-27T19:49:05.619Z" },
] ]
[package.optional-dependencies]
s3 = [
{ name = "mypy-boto3-s3" },
]
[[package]] [[package]]
name = "botocore" name = "botocore"
version = "1.40.12" version = "1.40.12"
@ -1871,7 +1876,7 @@ codegen = [
] ]
dev = [ dev = [
{ name = "black" }, { name = "black" },
{ name = "boto3-stubs" }, { name = "boto3-stubs", extra = ["s3"] },
{ name = "mypy" }, { name = "mypy" },
{ name = "nbval" }, { name = "nbval" },
{ name = "pandas-stubs" }, { name = "pandas-stubs" },
@ -1995,7 +2000,7 @@ codegen = [
] ]
dev = [ dev = [
{ name = "black" }, { name = "black" },
{ name = "boto3-stubs" }, { name = "boto3-stubs", extras = ["s3"] },
{ name = "mypy" }, { name = "mypy" },
{ name = "nbval" }, { name = "nbval" },
{ name = "pandas-stubs" }, { name = "pandas-stubs" },
@ -2568,6 +2573,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/87/e3/be76d87158ebafa0309946c4a73831974d4d6ab4f4ef40c3b53a385a66fd/mypy-1.18.2-py3-none-any.whl", hash = "sha256:22a1748707dd62b58d2ae53562ffc4d7f8bcc727e8ac7cbc69c053ddc874d47e", size = 2352367, upload-time = "2025-09-19T00:10:15.489Z" }, { url = "https://files.pythonhosted.org/packages/87/e3/be76d87158ebafa0309946c4a73831974d4d6ab4f4ef40c3b53a385a66fd/mypy-1.18.2-py3-none-any.whl", hash = "sha256:22a1748707dd62b58d2ae53562ffc4d7f8bcc727e8ac7cbc69c053ddc874d47e", size = 2352367, upload-time = "2025-09-19T00:10:15.489Z" },
] ]
[[package]]
name = "mypy-boto3-s3"
version = "1.40.26"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/00/b8/55d21ed9ca479df66d9892212ba7d7977850ef17aa80a83e3f11f31190fd/mypy_boto3_s3-1.40.26.tar.gz", hash = "sha256:8d2bfd1052894d0e84c9fb9358d838ba0eed0265076c7dd7f45622c770275c99", size = 75948, upload-time = "2025-09-08T20:12:21.405Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/85/a5/dba3384423834009bdd41c7021de5c663468a0e7bc4071cb301721e52a99/mypy_boto3_s3-1.40.26-py3-none-any.whl", hash = "sha256:6d055d16ef89a0133ade92f6b4f09603e4acc31a0f5e8f846edf4eb48f17b5a7", size = 82762, upload-time = "2025-09-08T20:12:19.338Z" },
]
[[package]] [[package]]
name = "mypy-extensions" name = "mypy-extensions"
version = "1.1.0" version = "1.1.0"