From 731ed6b6668e95d39718263f1fb8aea757036404 Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuelhwilliams@gmail.com>
Date: Thu, 29 Aug 2024 08:53:43 +0100
Subject: [PATCH] Fix type hints for boto3

---
 .pre-commit-config.yaml |  1 +
 data_store/aws.py       |  2 +-
 requirements-dev.in     |  2 +-
 requirements-dev.txt    | 12 ++++++++----
 tests/conftest.py       |  9 +++++----
 5 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e9ee9c193..267b3145c 100644
--- 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
index fc4e848dd..c836826f5 100644
--- 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
index a4b68927e..b7eca5338 100644
--- 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
index 95e1feb78..c5946d412 100644
--- 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
index 3c21fe6db..137ccac62 100644
--- 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)