Type response of Flask test client methods: get, post, etc. (#2904)

third_party/2and3/flask/app.pyi

@@ -127,7 +127,7 @@ class Flask(_PackageBoundObject):
     env: Optional[str] = ...
     debug: bool = ...
     def run(self, host: Optional[str] = ..., port: Optional[Union[int, str]] = ..., debug: Optional[bool] = ..., load_dotenv: bool = ..., **options: Any) -> None: ...
-    def test_client(self, use_cookies: bool = ..., **kwargs: Any) -> FlaskClient: ...
+    def test_client(self, use_cookies: bool = ..., **kwargs: Any) -> FlaskClient[Response]: ...
     def test_cli_runner(self, **kwargs: Any): ...
     def open_session(self, request: Any): ...
     def save_session(self, session: Any, response: Any): ...

third_party/2and3/flask/testing.pyi

@@ -4,17 +4,21 @@
 
 from click import BaseCommand
 from click.testing import CliRunner, Result
-from typing import Any, IO, Iterable, Mapping, Optional, Union
+from typing import Any, IO, Iterable, Mapping, Optional, Union, TypeVar
 from werkzeug.test import Client
 
 def make_test_environ_builder(app: Any, path: str = ..., base_url: Optional[Any] = ..., subdomain: Optional[Any] = ..., url_scheme: Optional[Any] = ..., *args: Any, **kwargs: Any): ...
 
-class FlaskClient(Client):
+# Response type for the client below.
+# By default _R is Tuple[Iterable[Any], Union[Text, int], werkzeug.datastructures.Headers], however
+# most commonly it is wrapped in a Reponse object.
+_R = TypeVar('_R')
+
+class FlaskClient(Client[_R]):
     preserve_context: bool = ...
     environ_base: Any = ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def session_transaction(self, *args: Any, **kwargs: Any) -> None: ...
-    def open(self, *args: Any, **kwargs: Any): ...
     def __enter__(self): ...
     def __exit__(self, exc_type: Any, exc_value: Any, tb: Any) -> None: ...
 

third_party/2and3/werkzeug/test.pyi

@@ -1,5 +1,7 @@
 import sys
-from typing import Any, Optional, Text
+from wsgiref.types import WSGIEnvironment
+from typing import Any, Generic, Optional, Text, Tuple, Type, TypeVar, overload
+from typing_extensions import Literal
 
 if sys.version_info < (3,):
     from urllib2 import Request as U2Request
@@ -67,12 +69,16 @@ class EnvironBuilder:
 
 class ClientRedirectError(Exception): ...
 
-class Client:
+# Response type for the client below.
+# By default _R is Tuple[Iterable[Any], Union[Text, int], datastructures.Headers]
+_R = TypeVar('_R')
+
+class Client(Generic[_R]):
     application: Any
-    response_wrapper: Any
+    response_wrapper: Optional[Type[_R]]
     cookie_jar: Any
     allow_subdomain_redirects: Any
-    def __init__(self, application, response_wrapper: Optional[Any] = ..., use_cookies: bool = ...,
+    def __init__(self, application, response_wrapper: Optional[Type[_R]] = ..., use_cookies: bool = ...,
                  allow_subdomain_redirects: bool = ...): ...
     def set_cookie(self, server_name, key, value: str = ..., max_age: Optional[Any] = ..., expires: Optional[Any] = ...,
                    path: str = ..., domain: Optional[Any] = ..., secure: Optional[Any] = ..., httponly: bool = ...,
@@ -80,15 +86,69 @@ class Client:
     def delete_cookie(self, server_name, key, path: str = ..., domain: Optional[Any] = ...): ...
     def run_wsgi_app(self, environ, buffered: bool = ...): ...
     def resolve_redirect(self, response, new_location, environ, buffered: bool = ...): ...
-    def open(self, *args, **kwargs): ...
-    def get(self, *args, **kw): ...
-    def patch(self, *args, **kw): ...
-    def post(self, *args, **kw): ...
-    def head(self, *args, **kw): ...
-    def put(self, *args, **kw): ...
-    def delete(self, *args, **kw): ...
-    def options(self, *args, **kw): ...
-    def trace(self, *args, **kw): ...
+
+    @overload
+    def open(self, *args, as_tuple: Literal[True], **kwargs) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def open(self, *args, as_tuple: Literal[False] = ..., **kwargs) -> _R: ...
+    @overload
+    def open(self, *args, as_tuple: bool, **kwargs) -> Any: ...
+
+    @overload
+    def get(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def get(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def get(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def patch(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def patch(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def patch(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def post(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def post(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def post(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def head(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def head(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def head(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def put(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def put(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def put(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def delete(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def delete(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def delete(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def options(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def options(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def options(self, *args, as_tuple: bool, **kw) -> Any: ...
+
+    @overload
+    def trace(self, *args, as_tuple: Literal[True], **kw) -> Tuple[WSGIEnvironment, _R]: ...
+    @overload
+    def trace(self, *args, as_tuple: Literal[False] = ..., **kw) -> _R: ...
+    @overload
+    def trace(self, *args, as_tuple: bool, **kw) -> Any: ...
 
 def create_environ(*args, **kwargs): ...
 def run_wsgi_app(app, environ, buffered: bool = ...): ...