mypy: Add boto3-stubs.

Signed-off-by: Anders Kaseorg <anders@zulip.com>

pyproject.toml

@@ -51,8 +51,6 @@ module = [
     "aioapns.*",
     "bitfield.*",
     "bmemcached.*",
-    "boto3.*",
-    "botocore.*",
     "bs4.*",
     "bson.*",
     "cairosvg.*",

requirements/common.in

@@ -36,6 +36,7 @@ SQLAlchemy==1.3.*  # 1.4 has badly busted type annotations
 
 # Needed for S3 file uploads
 boto3
+mypy-boto3-s3
 
 # Needed for integrations
 defusedxml

requirements/dev.txt

@@ -101,6 +101,10 @@ boto3==1.17.105 \
     # via
     #   -r requirements/common.in
     #   moto
+boto3-stubs[s3]==1.18.17 \
+    --hash=sha256:97aef3a2173bedd95b75aafaa1a6a85e321af107a11855f30afded7a1e2462bc \
+    --hash=sha256:b5bd2f3f54f06eecb9f6a085643c4c02a057c5fbbad4256c4e4a02a0744758df
+    # via -r requirements/mypy.in
 botocore==1.20.105 \
     --hash=sha256:b0fda4edf8eb105453890700d49011ada576d0cc7326a0699dfabe9e872f552c \
     --hash=sha256:b5ba72d22212b0355f339c2a98b3296b3b2202a48e6a2b1366e866bc65a64b67
@@ -108,6 +112,10 @@ botocore==1.20.105 \
     #   boto3
     #   moto
     #   s3transfer
+botocore-stubs==1.21.17 \
+    --hash=sha256:6fca2ff326532e8ad8b74c1e5ef6e0457f409ebe38cb8b4aa6cb50b534ee3ee3 \
+    --hash=sha256:b754cb23471948b8cbe50d24d4a74c617672905e4ae170b01bcfcae6496d7cb6
+    # via boto3-stubs
 cachetools==4.2.2 \
     --hash=sha256:2cc0b89715337ab6dbba85b5b50effe2b0c74e035d83ee8ed637cf52f12ae001 \
     --hash=sha256:61b5ed1e22a0924aed1d23b478f37e8d52549ff8a961de2909c69bf950020cff
@@ -788,6 +796,12 @@ mypy==0.910 \
     # via
     #   -r requirements/mypy.in
     #   sqlalchemy-stubs
+mypy-boto3-s3==1.18.17 \
+    --hash=sha256:63a76e94df730984196fd46be3f541dacc8d162f6f70c210a4cd9a80d6775e3b \
+    --hash=sha256:af3699fb37614ff8044b7b6d3d7dd2211e5307bf018ac4f0a3591ec2011123c1
+    # via
+    #   -r requirements/common.in
+    #   boto3-stubs
 mypy-extensions==0.4.3 \
     --hash=sha256:090fedd75945a69ae91ce1303b5824f428daf5a028d2f6ab8a299250a846f15d \
     --hash=sha256:2d82818f5bb3e369420cb3c4060a7970edba416647068eb4c5343488a6c604a8
@@ -1742,9 +1756,12 @@ typing-extensions==3.10.0.0 \
     #   arrow
     #   asgiref
     #   black
+    #   boto3-stubs
+    #   botocore-stubs
     #   importlib-metadata
     #   libcst
     #   mypy
+    #   mypy-boto3-s3
     #   pyre-check
     #   pyre-extensions
     #   sqlalchemy-stubs

requirements/mypy.in

@@ -5,6 +5,7 @@
 mypy
 
 backoff-stubs
+boto3-stubs[s3]
 lxml-stubs
 https://github.com/andersk/pika-stubs/archive/87c5795741449e37bdbd2ceceee853fd56462440.zip#egg=pika-stubs==0.1.3+git  # https://github.com/hahow/pika-stubs/issues/1, https://github.com/hahow/pika-stubs/pull/4
 sqlalchemy-stubs

requirements/mypy.txt

@@ -10,6 +10,14 @@
 backoff-stubs==1.10.0 \
     --hash=sha256:03e995de0a70016c6fe758498e1ca811f1db517c00cbd06e3039c9e4f6ea2566
     # via -r requirements/mypy.in
+boto3-stubs[s3]==1.18.17 \
+    --hash=sha256:97aef3a2173bedd95b75aafaa1a6a85e321af107a11855f30afded7a1e2462bc \
+    --hash=sha256:b5bd2f3f54f06eecb9f6a085643c4c02a057c5fbbad4256c4e4a02a0744758df
+    # via -r requirements/mypy.in
+botocore-stubs==1.21.17 \
+    --hash=sha256:6fca2ff326532e8ad8b74c1e5ef6e0457f409ebe38cb8b4aa6cb50b534ee3ee3 \
+    --hash=sha256:b754cb23471948b8cbe50d24d4a74c617672905e4ae170b01bcfcae6496d7cb6
+    # via boto3-stubs
 lxml-stubs==0.2.0 \
     --hash=sha256:78f1bfb31b1f2af9a5c9e9a602ab1b589a64a5a3cc444931a39cdfd02d6864b0 \
     --hash=sha256:f0b3621ec2a23bea4145f484490c8b27383ecb407b3f8b079199ad4a0af4180b
@@ -41,6 +49,10 @@ mypy==0.910 \
     # via
     #   -r requirements/mypy.in
     #   sqlalchemy-stubs
+mypy-boto3-s3==1.18.17 \
+    --hash=sha256:63a76e94df730984196fd46be3f541dacc8d162f6f70c210a4cd9a80d6775e3b \
+    --hash=sha256:af3699fb37614ff8044b7b6d3d7dd2211e5307bf018ac4f0a3591ec2011123c1
+    # via boto3-stubs
 mypy-extensions==0.4.3 \
     --hash=sha256:090fedd75945a69ae91ce1303b5824f428daf5a028d2f6ab8a299250a846f15d \
     --hash=sha256:2d82818f5bb3e369420cb3c4060a7970edba416647068eb4c5343488a6c604a8
@@ -179,5 +191,6 @@ typing-extensions==3.10.0.0 \
     --hash=sha256:50b6f157849174217d0656f99dc82fe932884fb250826c18350e159ec6cdf342 \
     --hash=sha256:779383f6086d90c99ae41cf0ff39aac8a7937a9283ce0a414e5dd782f4c94a84
     # via
+    #   boto3-stubs
     #   mypy
     #   sqlalchemy-stubs

requirements/prod.txt

@@ -494,6 +494,10 @@ more-itertools==8.8.0 \
     --hash=sha256:2cf89ec599962f2ddc4d568a05defc40e0a587fbc10d5989713638864c36be4d \
     --hash=sha256:83f0308e05477c68f56ea3a888172c78ed5d5b3c282addb67508e7ba6c8f813a
     # via openapi-core
+mypy-boto3-s3==1.18.17 \
+    --hash=sha256:63a76e94df730984196fd46be3f541dacc8d162f6f70c210a4cd9a80d6775e3b \
+    --hash=sha256:af3699fb37614ff8044b7b6d3d7dd2211e5307bf018ac4f0a3591ec2011123c1
+    # via -r requirements/common.in
 oauthlib==3.1.1 \
     --hash=sha256:42bf6354c2ed8c6acb54d971fce6f88193d97297e18602a3a886603f9d7730cc \
     --hash=sha256:8f0215fcc533dd8dd1bee6f4c412d4f0cd7297307d43ac61666389e3bc3198a3
@@ -1046,6 +1050,7 @@ typing-extensions==3.10.0.0 \
     #   -r requirements/common.in
     #   asgiref
     #   importlib-metadata
+    #   mypy-boto3-s3
     #   zulip-bots
 uhashring==2.1 \
     --hash=sha256:b21340d0d32497a67f34f5177a64908115fdc23264ed87fa7d1eca79ef9641fa

version.py

@@ -48,4 +48,4 @@ API_FEATURE_LEVEL = 92
 #   historical commits sharing the same major version, in which case a
 #   minor version bump suffices.
 
-PROVISION_VERSION = "153.13"
+PROVISION_VERSION = "153.14"

zerver/lib/export.py

@@ -16,12 +16,12 @@ import tempfile
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import orjson
-from boto3.resources.base import ServiceResource
 from django.apps import apps
 from django.conf import settings
 from django.forms.models import model_to_dict
 from django.utils.timezone import is_naive as timezone_is_naive
 from django.utils.timezone import make_aware as timezone_make_aware
+from mypy_boto3_s3.service_resource import Object
 
 import zerver.lib.upload
 from analytics.models import RealmCount, StreamCount, UserCount
@@ -1275,7 +1275,7 @@ def export_uploads_and_avatars(realm: Realm, output_dir: Path) -> None:
 
 def _check_key_metadata(
     email_gateway_bot: Optional[UserProfile],
-    key: ServiceResource,
+    key: Object,
     processing_avatars: bool,
     realm: Realm,
     user_ids: Set[int],
@@ -1298,10 +1298,10 @@ def _check_key_metadata(
 
 
 def _get_exported_s3_record(
-    bucket_name: str, key: ServiceResource, processing_emoji: bool
+    bucket_name: str, key: Object, processing_emoji: bool
-) -> Dict[str, Union[str, int]]:
+) -> Dict[str, Any]:
     # Helper function for export_files_from_s3
-    record = dict(
+    record: Dict[str, Any] = dict(
         s3_path=key.key,
         bucket=bucket_name,
         size=key.content_length,
@@ -1315,7 +1315,7 @@ def _get_exported_s3_record(
         record["file_name"] = os.path.basename(key.key)
 
     if "user_profile_id" in record:
-        user_profile = get_user_profile_by_id(record["user_profile_id"])
+        user_profile = get_user_profile_by_id(int(record["user_profile_id"]))
         record["user_profile_email"] = user_profile.email
 
         # Fix the record ids
@@ -1340,7 +1340,7 @@ def _get_exported_s3_record(
 
 
 def _save_s3_object_to_file(
-    key: ServiceResource,
+    key: Object,
     output_dir: str,
     processing_avatars: bool,
     processing_emoji: bool,
@@ -1365,7 +1365,7 @@ def _save_s3_object_to_file(
 
     if not os.path.exists(dirname):
         os.makedirs(dirname)
-    key.download_file(filename)
+    key.download_file(Filename=filename)
 
 
 def export_files_from_s3(

zerver/lib/import_realm.py

@@ -823,7 +823,7 @@ def import_uploads(
                     content_type = "application/octet-stream"
 
             key.upload_file(
-                os.path.join(import_dir, record["path"]),
+                Filename=os.path.join(import_dir, record["path"]),
                 ExtraArgs={"ContentType": content_type, "Metadata": metadata},
             )
         else:

zerver/lib/test_helpers.py

@@ -28,7 +28,6 @@ import boto3
 import fakeldap
 import ldap
 import orjson
-from boto3.resources.base import ServiceResource
 from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
 from django.db.migrations.state import StateApps
@@ -37,6 +36,7 @@ from django.http.request import QueryDict
 from django.test import override_settings
 from django.urls import URLResolver
 from moto import mock_s3
+from mypy_boto3_s3.service_resource import Bucket
 
 import zerver.lib.upload
 from zerver.lib import cache
@@ -521,7 +521,7 @@ def use_s3_backend(method: FuncT) -> FuncT:
     return new_method
 
 
-def create_s3_buckets(*bucket_names: str) -> List[ServiceResource]:
+def create_s3_buckets(*bucket_names: str) -> List[Bucket]:
     session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
     s3 = session.resource("s3")
     buckets = [s3.create_bucket(Bucket=name) for name in bucket_names]

zerver/lib/upload.py

@@ -15,7 +15,6 @@ from typing import Any, Callable, Optional, Tuple
 
 import boto3
 import botocore
-from boto3.resources.base import ServiceResource
 from boto3.session import Session
 from botocore.client import Config
 from django.conf import settings
@@ -25,6 +24,8 @@ from django.http import HttpRequest
 from django.urls import reverse
 from django.utils.translation import gettext as _
 from jinja2.utils import Markup as mark_safe
+from mypy_boto3_s3.client import S3Client
+from mypy_boto3_s3.service_resource import Bucket, Object
 from PIL import Image, ImageOps
 from PIL.GifImagePlugin import GifImageFile
 from PIL.Image import DecompressionBombError
@@ -278,9 +279,7 @@ class ZulipUploadBackend:
 ### S3
 
 
-def get_bucket(bucket_name: str, session: Optional[Session] = None) -> ServiceResource:
+def get_bucket(bucket_name: str, session: Optional[Session] = None) -> Bucket:
-    # See https://github.com/python/typeshed/issues/2706
-    # for why this return type is a `ServiceResource`.
     if session is None:
         session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
     bucket = session.resource(
@@ -290,8 +289,7 @@ def get_bucket(bucket_name: str, session: Optional[Session] = None) -> ServiceRe
 
 
 def upload_image_to_s3(
-    # See https://github.com/python/typeshed/issues/2706
+    bucket: Bucket,
-    bucket: ServiceResource,
     file_name: str,
     content_type: Optional[str],
     user_profile: UserProfile,
@@ -367,7 +365,7 @@ class S3UploadBackend(ZulipUploadBackend):
         self.avatar_bucket = get_bucket(settings.S3_AVATAR_BUCKET, self.session)
         self.uploads_bucket = get_bucket(settings.S3_AUTH_UPLOADS_BUCKET, self.session)
 
-        self._boto_client = None
+        self._boto_client: Optional[S3Client] = None
         self.public_upload_url_base = self.construct_public_upload_url_base()
 
     def construct_public_upload_url_base(self) -> str:
@@ -410,7 +408,7 @@ class S3UploadBackend(ZulipUploadBackend):
         assert not key.startswith("/")
         return urllib.parse.urljoin(self.public_upload_url_base, key)
 
-    def get_boto_client(self) -> botocore.client.BaseClient:
+    def get_boto_client(self) -> S3Client:
         """
         Creating the client takes a long time so we need to cache it.
         """
@@ -424,7 +422,7 @@ class S3UploadBackend(ZulipUploadBackend):
             )
         return self._boto_client
 
-    def delete_file_from_s3(self, path_id: str, bucket: ServiceResource) -> bool:
+    def delete_file_from_s3(self, path_id: str, bucket: Bucket) -> bool:
         key = bucket.Object(path_id)
 
         try:
@@ -532,9 +530,7 @@ class S3UploadBackend(ZulipUploadBackend):
         self.delete_file_from_s3(path_id + "-medium.png", self.avatar_bucket)
         self.delete_file_from_s3(path_id, self.avatar_bucket)
 
-    def get_avatar_key(self, file_name: str) -> ServiceResource:
+    def get_avatar_key(self, file_name: str) -> Object:
-        # See https://github.com/python/typeshed/issues/2706
-        # for why this return type is a `ServiceResource`.
         key = self.avatar_bucket.Object(file_name)
         return key
 
@@ -693,7 +689,10 @@ class S3UploadBackend(ZulipUploadBackend):
             os.path.join("exports", secrets.token_hex(16), os.path.basename(tarball_path))
         )
 
-        key.upload_file(tarball_path, Callback=percent_callback)
+        if percent_callback is None:
+            key.upload_file(Filename=tarball_path)
+        else:
+            key.upload_file(Filename=tarball_path, Callback=percent_callback)
 
         public_url = self.get_public_upload_url(key.key)
         return public_url

zerver/migrations/0149_realm_emoji_drop_unique_constraint.py

@@ -6,6 +6,7 @@ from django.conf import settings
 from django.db import migrations, models
 from django.db.backends.postgresql.schema import DatabaseSchemaEditor
 from django.db.migrations.state import StateApps
+from mypy_boto3_s3.type_defs import CopySourceTypeDef
 
 
 class Uploader:
@@ -66,8 +67,8 @@ class S3Uploader(Uploader):
         ).Bucket(self.bucket_name)
 
     def copy_files(self, src_key: str, dst_key: str) -> None:
-        source = dict(Bucket=self.bucket_name, Key=src_key)
+        source = CopySourceTypeDef(Bucket=self.bucket_name, Key=src_key)
-        self.bucket.copy(source, dst_key)
+        self.bucket.copy(CopySource=source, Key=dst_key)
 
 
 def get_uploader() -> Uploader: