Fix type hints for boto3

communitiesuk · Aug 29, 2024 · 731ed6b · 731ed6b
1 parent 633a805
commit 731ed6b
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 10 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -26,6 +26,7 @@ repos:
       # args: [--show-error-codes]
       exclude: ^(submit|find)/.*$
       additional_dependencies: [
+        "boto3-stubs[s3,logs]",
         pandas-stubs,
         sqlalchemy-stubs,
         types-beautifulsoup4,

diff --git a/data_store/aws.py b/data_store/aws.py
@@ -49,7 +49,7 @@ def upload_file(file: Union[IO, FileStorage], bucket: str, object_name: str, met
 
     file.seek(0)
     _S3_CLIENT.upload_fileobj(
-        file,
+        file.stream if isinstance(file, FileStorage) else file,
         bucket,
         object_name,
         ExtraArgs={"Metadata": metadata if metadata else {}, "ContentType": content_type},

diff --git a/requirements-dev.in b/requirements-dev.in
@@ -17,10 +17,10 @@ pytest-playwright
 #   Code Quality
 #-----------------------------------
 mypy==1.11.1  # If bumping this, please also bump .pre-commit-config.yml
+boto3-stubs[s3,logs]
 pandas-stubs
 sqlalchemy-stubs
 types-beautifulsoup4
-types-boto3
 types-colorama
 types-flask
 types-Flask-Migrate

diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -34,8 +34,8 @@ blinker==1.8.2
     #   sentry-sdk
 boto3==1.34.151
     # via -r requirements.txt
-boto3-stubs==1.34.133
-    # via types-boto3
+boto3-stubs[logs,s3]==1.34.133
+    # via -r requirements-dev.in
 botocore==1.34.151
     # via
     #   -r requirements.txt
@@ -231,6 +231,10 @@ mypy==1.11.1
     # via
     #   -r requirements-dev.in
     #   sqlalchemy-stubs
+mypy-boto3-logs==1.34.151
+    # via boto3-stubs
+mypy-boto3-s3==1.34.162
+    # via boto3-stubs
 mypy-extensions==1.0.0
     # via
     #   -r requirements.txt
@@ -415,8 +419,6 @@ types-awscrt==0.20.12
     # via botocore-stubs
 types-beautifulsoup4==4.12.0.20240511
     # via -r requirements-dev.in
-types-boto3==1.0.2
-    # via -r requirements-dev.in
 types-click==7.1.8
     # via types-flask
 types-colorama==0.4.15.20240311
@@ -473,6 +475,8 @@ typing-extensions==4.6.0
     #   alembic
     #   boto3-stubs
     #   mypy
+    #   mypy-boto3-logs
+    #   mypy-boto3-s3
     #   pydantic
     #   pyee
     #   sqlalchemy

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -19,6 +19,7 @@
 import pandas as pd
 import pytest
 from flask.testing import FlaskClient
+from mypy_boto3_s3.type_defs import ObjectIdentifierTypeDef
 from sqlalchemy import text
 from werkzeug.test import TestResponse
 
@@ -731,10 +732,10 @@ def create_bucket(bucket: str):
 
 def delete_bucket(bucket: str):
     """Helper function that deletes all objects in a specified bucket and then deletes the bucket."""
-    objects = _S3_CLIENT.list_objects_v2(Bucket=bucket)
-    if objects := objects.get("Contents"):
-        objects = list(map(lambda x: {"Key": x["Key"]}, objects))
-        _S3_CLIENT.delete_objects(Bucket=bucket, Delete={"Objects": objects})
+    s3_objects_response = _S3_CLIENT.list_objects_v2(Bucket=bucket)
+    if "Contents" in s3_objects_response:
+        s3_objects: list[ObjectIdentifierTypeDef] = [{"Key": obj["Key"]} for obj in s3_objects_response["Contents"]]
+        _S3_CLIENT.delete_objects(Bucket=bucket, Delete={"Objects": s3_objects})
     _S3_CLIENT.delete_bucket(Bucket=bucket)