Improve the decorator typing

This should make the type checking of decorated functions better and
hence more likely to catch bugs.

Note that the validation decorators cannot yet be improved as they
inject keyword arguments, which is not supported as per
https://peps.python.org/pep-0612/#concatenating-keyword-parameters .

Author: pgjones

src/quart_schema/documentation.py

@@ -1,6 +1,4 @@
-from typing import Callable, Dict, Optional, Tuple
-
-from quart import ResponseReturnValue as QuartResponseReturnValue
+from typing import Callable, Dict, Optional, Tuple, TypeVar
 
 from .typing import Model
 from .validation import (
@@ -11,8 +9,10 @@
     QUART_SCHEMA_RESPONSE_ATTRIBUTE,
 )
 
+T = TypeVar("T", bound=Callable)
+
 
-def document_querystring(model_class: Model) -> Callable:
+def document_querystring(model_class: Model) -> Callable[[T], T]:
     """Document the request querystring arguments.
 
     Add the request querystring **model_class** to the openapi
@@ -25,15 +25,15 @@ def document_querystring(model_class: Model) -> Callable:
 
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         setattr(func, QUART_SCHEMA_QUERYSTRING_ATTRIBUTE, model_class)
 
         return func
 
     return decorator
 
 
-def document_headers(model_class: Model) -> Callable:
+def document_headers(model_class: Model) -> Callable[[T], T]:
     """Document the request headers.
 
     Add the request **model_class** to the openapi generated
@@ -46,7 +46,7 @@ def document_headers(model_class: Model) -> Callable:
 
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         setattr(func, QUART_SCHEMA_HEADERS_ATTRIBUTE, model_class)
 
         return func
@@ -58,7 +58,7 @@ def document_request(
     model_class: Model,
     *,
     source: DataSource = DataSource.JSON,
-) -> Callable:
+) -> Callable[[T], T]:
     """Document the request data.
 
     Add the request **model_class** to the openapi generated
@@ -72,7 +72,7 @@ def document_request(
             encoded).
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         setattr(func, QUART_SCHEMA_REQUEST_ATTRIBUTE, (model_class, source))
 
         return func
@@ -84,7 +84,7 @@ def document_response(
     model_class: Model,
     status_code: int = 200,
     headers_model_class: Optional[Model] = None,
-) -> Callable:
+) -> Callable[[T], T]:
     """Document the response data.
 
     Add the response **model_class**, and its corresponding (optional)
@@ -103,9 +103,7 @@ def document_response(
 
     """
 
-    def decorator(
-        func: Callable[..., QuartResponseReturnValue]
-    ) -> Callable[..., QuartResponseReturnValue]:
+    def decorator(func: T) -> T:
         schemas = getattr(func, QUART_SCHEMA_RESPONSE_ATTRIBUTE, {})
         schemas[status_code] = (model_class, headers_model_class)
         setattr(func, QUART_SCHEMA_RESPONSE_ATTRIBUTE, schemas)
@@ -122,15 +120,15 @@ def document(
     request_source: DataSource = DataSource.JSON,
     headers: Optional[Model] = None,
     responses: Dict[int, Tuple[Model, Optional[Model]]],
-) -> Callable:
+) -> Callable[[T], T]:
     """Document the route.
 
     This is a shorthand combination of of the document_querystring,
     document_request, document_headers, and document_response
     decorators. Please see the docstrings for those decorators.
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         if querystring is not None:
             func = document_querystring(querystring)(func)
         if request is not None:

src/quart_schema/extension.py

@@ -5,7 +5,7 @@
 from collections import defaultdict
 from functools import wraps
 from types import new_class
-from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple, TypeVar, Union
 
 import click
 import humps
@@ -48,6 +48,7 @@
 except ImportError:
     to_builtins = None
 
+T = TypeVar("T", bound=Callable)
 
 SecurityScheme = Union[
     APIKeySecurityScheme,
@@ -150,7 +151,7 @@ def default(object_: Any) -> Any:
     return JSONProvider(app)
 
 
-def hide(func: Callable) -> Callable:
+def hide(func: T) -> T:
     """Mark the func as hidden.
 
     This will prevent the route from being included in the
@@ -398,7 +399,7 @@ async def decorator(result: ResponseReturnValue) -> QuartResponseReturnValue:
     return decorator
 
 
-def operation_id(operationid: str) -> Callable:
+def operation_id(operationid: str) -> Callable[[T], T]:
     """Override the operationId of the route.
 
     This allows for overriding the operationId, which is normally calculated from the
@@ -409,15 +410,15 @@ def operation_id(operationid: str) -> Callable:
 
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         setattr(func, QUART_SCHEMA_OPERATION_ID_ATTRIBUTE, str(operationid))
 
         return func
 
     return decorator
 
 
-def tag(tags: Iterable[str]) -> Callable:
+def tag(tags: Iterable[str]) -> Callable[[T], T]:
     """Add tag names to the route.
 
     This allows for tags to be associated with the route, thereby
@@ -428,22 +429,22 @@ def tag(tags: Iterable[str]) -> Callable:
 
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         setattr(func, QUART_SCHEMA_TAG_ATTRIBUTE, set(tags))
 
         return func
 
     return decorator
 
 
-def deprecate(func: Callable) -> Callable:
+def deprecate(func: T) -> T:
     """Mark endpoint as deprecated."""
     setattr(func, QUART_SCHEMA_DEPRECATED, True)
 
     return func
 
 
-def security_scheme(schemes: Iterable[Dict[str, List[str]]]) -> Callable:
+def security_scheme(schemes: Iterable[Dict[str, List[str]]]) -> Callable[[T], T]:
     """Add security schemes to the route.
 
     Allows security schemes to be associated with this route. Security
@@ -455,7 +456,7 @@ def security_scheme(schemes: Iterable[Dict[str, List[str]]]) -> Callable:
 
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: T) -> T:
         setattr(func, QUART_SCHEMA_SECURITY_ATTRIBUTE, schemes)
 
         return func