Add mypy annotations and mark it as checkable.

.pre-commit-config.yaml

@@ -44,6 +44,13 @@ repos:
         additional_dependencies:
           - "flake8-bugbear==24.8.19"
 
+  - repo: "https://github.com/pre-commit/mirrors-mypy"
+    rev: "v1.19.1"
+    hooks:
+      - id: "mypy"
+        additional_dependencies:
+          - "flask>=3.0"
+
   - repo: "https://github.com/python-jsonschema/check-jsonschema"
     rev: "0.36.0"
     hooks:

flask_compress/__init__.py

@@ -1,4 +1,4 @@
-from .flask_compress import Compress, DictCache
+from .flask_compress import CacheBackend, Compress, DictCache
 
 # _version.py is generated by setuptools_scm when building the package.
 # It is not version-controlled, so if it is missing, this likely means that
@@ -9,4 +9,4 @@
     __version__ = "0"
 
 
-__all__ = ("Compress", "DictCache")
+__all__ = ("CacheBackend", "Compress", "DictCache")

flask_compress/flask_compress.py

@@ -2,34 +2,44 @@
 # Copyright (c) 2013-2017 William Fagan
 # License: The MIT License (MIT)
 
+from __future__ import annotations
+
 import functools
 from collections import defaultdict
+from collections.abc import Callable, Iterator
 from functools import lru_cache
+from typing import Any, Protocol
 
 try:
     import brotlicffi as brotli
 except ImportError:
     import brotli
 
-from flask import after_this_request, current_app, request, stream_with_context
+from flask import Flask, after_this_request, current_app, request, stream_with_context
+from flask.wrappers import Response
 
 from .compat import compression
 
 
+class CacheBackend(Protocol):
+    def get(self, key: str) -> bytes | None: ...
+    def set(self, key: str, value: bytes) -> bool | None: ...
+
+
 class DictCache:
 
-    def __init__(self):
-        self.data = {}
+    def __init__(self) -> None:
+        self.data: dict[str, bytes] = {}
 
-    def get(self, key):
+    def get(self, key: str) -> bytes | None:
         return self.data.get(key)
 
-    def set(self, key, value):
+    def set(self, key: str, value: bytes) -> None:
         self.data[key] = value
 
 
 @lru_cache(maxsize=128)
-def _choose_algorithm(algorithms, accept_encoding):
+def _choose_algorithm(algorithms: tuple[str, ...], accept_encoding: str) -> str | None:
     """
     Determine which compression algorithm we're going to use based on the
     client request. The `Accept-Encoding` header may list one or more desired
@@ -46,7 +56,7 @@ def _choose_algorithm(algorithms, accept_encoding):
     fallback_to_any = False
 
     # Map quality factors to requested algorithm names.
-    algos_by_quality = defaultdict(set)
+    algos_by_quality: defaultdict[float, set[str | None]] = defaultdict(set)
 
     # Set of supported algorithms
     server_algos_set = set(algorithms)
@@ -96,7 +106,7 @@ def _choose_algorithm(algorithms, accept_encoding):
     return None
 
 
-def _format(algo):
+def _format(algo: str | list[str]) -> tuple[str, ...]:
     """Format the algorithm configuration into a tuple of strings.
 
     >>> _format("gzip, deflate, br")
@@ -122,7 +132,14 @@ class Compress:
     :type app: :class:`flask.Flask` or None
     """
 
-    def __init__(self, app=None):
+    cache: CacheBackend | None
+    cache_key: Callable[..., str] | None
+    compress_mimetypes_set: set[str]
+    enabled_algorithms: tuple[str, ...]
+    streaming_algorithms: tuple[str, ...]
+    streaming_endpoint_with_conditional: set[str]
+
+    def __init__(self, app: Flask | None = None) -> None:
         """
         An alternative way to pass your :class:`flask.Flask` application
         object to Flask-Compress. :meth:`init_app` also takes care of some
@@ -134,7 +151,7 @@ def __init__(self, app=None):
         if app is not None:
             self.init_app(app)
 
-    def init_app(self, app):
+    def init_app(self, app: Flask) -> None:
         defaults = [
             (
                 "COMPRESS_MIMETYPES",
@@ -202,7 +219,7 @@ def init_app(self, app):
         if app.config["COMPRESS_REGISTER"] and app.config["COMPRESS_MIMETYPES"]:
             app.after_request(self.after_request)
 
-    def after_request(self, response):
+    def after_request(self, response: Response) -> Response:
         app = self.app or current_app
 
         vary = response.headers.get("Vary")
@@ -247,6 +264,7 @@ def after_request(self, response):
             response.headers.pop("Content-Length", None)
         else:
             if self.cache is not None:
+                assert self.cache_key is not None
                 key = f"{chosen_algorithm};{self.cache_key(request)}"
                 compressed_content = self.cache.get(key)
                 if compressed_content is None:
@@ -265,7 +283,7 @@ def after_request(self, response):
         etag, is_weak = response.get_etag()
 
         if etag and not is_weak:
-            response.set_etag(f"{etag}:{chosen_algorithm}", weak=is_weak)
+            response.set_etag(f"{etag}:{chosen_algorithm}", weak=False)
 
         if (
             app.config["COMPRESS_EVALUATE_CONDITIONAL_REQUEST"]
@@ -276,12 +294,12 @@ def after_request(self, response):
 
         return response
 
-    def compressed(self):
-        def decorator(f):
+    def compressed(self) -> Callable[..., Callable[..., Any]]:
+        def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
             @functools.wraps(f)
-            def decorated_function(*args, **kwargs):
+            def decorated_function(*args: Any, **kwargs: Any) -> Any:
                 @after_this_request
-                def compressor(response):
+                def compressor(response: Response) -> Response:
                     return self.after_request(response)
 
                 return f(*args, **kwargs)
@@ -291,18 +309,24 @@ def compressor(response):
         return decorator
 
 
-def _compress_data(app, data, algorithm):
+def _compress_data(app: Flask, data: bytes, algorithm: str) -> bytes:
     if algorithm == "zstd":
-        return compression.zstd.compress(data, app.config["COMPRESS_ZSTD_LEVEL"])
+        return compression.zstd.compress(  # type: ignore[no-any-return]
+            data, app.config["COMPRESS_ZSTD_LEVEL"]
+        )
 
     if algorithm == "gzip":
-        return compression.gzip.compress(data, app.config["COMPRESS_LEVEL"])
+        return compression.gzip.compress(  # type: ignore[no-any-return]
+            data, app.config["COMPRESS_LEVEL"]
+        )
 
     if algorithm == "deflate":
-        return compression.zlib.compress(data, app.config["COMPRESS_DEFLATE_LEVEL"])
+        return compression.zlib.compress(  # type: ignore[no-any-return]
+            data, app.config["COMPRESS_DEFLATE_LEVEL"]
+        )
 
     if algorithm == "br":
-        return brotli.compress(
+        return brotli.compress(  # type: ignore[no-any-return]
             data,
             mode=app.config["COMPRESS_BR_MODE"],
             quality=app.config["COMPRESS_BR_LEVEL"],
@@ -313,21 +337,23 @@ def _compress_data(app, data, algorithm):
     raise ValueError(f"Unknown compression algorithm: {algorithm}")
 
 
-def _uncompress_data(data, algorithm):
+def _uncompress_data(data: bytes, algorithm: str) -> bytes:
     # This is used for tests purposes only.
     if algorithm == "zstd":
-        return compression.zstd.decompress(data)
+        return compression.zstd.decompress(data)  # type: ignore[no-any-return]
     if algorithm == "gzip":
-        return compression.gzip.decompress(data)
+        return compression.gzip.decompress(data)  # type: ignore[no-any-return]
     if algorithm == "deflate":
-        return compression.zlib.decompress(data)
+        return compression.zlib.decompress(data)  # type: ignore[no-any-return]
     if algorithm == "br":
-        return brotli.decompress(data)
+        return brotli.decompress(data)  # type: ignore[no-any-return]
 
     raise ValueError(f"Unknown compression algorithm: {algorithm}")
 
 
-def _compress_chunks(app, chunks, algorithm):
+def _compress_chunks(
+    app: Flask, chunks: Iterator[bytes], algorithm: str
+) -> Iterator[bytes]:
     if algorithm == "zstd":
         level = app.config["COMPRESS_ZSTD_LEVEL"]
         compressor = compression.zstd.ZstdCompressor(level=level)

pyproject.toml

@@ -33,3 +33,14 @@ fail_under = 94
 [tool.coverage.html]
 directory = "htmlcov/"
 skip_covered = false
+
+
+# type checking
+# --------------
+
+[tool.mypy]
+strict = true
+
+[[tool.mypy.overrides]]
+module = ["brotli", "brotlicffi", "backports.*", "flask_caching"]
+ignore_missing_imports = true

setup.py

@@ -16,6 +16,7 @@
     long_description=LONG_DESCRIPTION,
     long_description_content_type="text/markdown",
     packages=find_packages(exclude=["tests"]),
+    package_data={"flask_compress": ["py.typed"]},
     include_package_data=True,
     platforms="any",
     python_requires=">=3.9",