diff --git a/chalice/local.py b/chalice/local.py index f6ce45102f..3214f8de3d 100644 --- a/chalice/local.py +++ b/chalice/local.py @@ -16,7 +16,16 @@ from six.moves.BaseHTTPServer import HTTPServer from six.moves.BaseHTTPServer import BaseHTTPRequestHandler from six.moves.socketserver import ThreadingMixIn -from typing import List, Any, Dict, Tuple, Callable, Optional, Union # noqa +from typing import ( + List, + Any, + Dict, + Tuple, + Callable, + Optional, + Union, + cast, +) # noqa from chalice.app import Chalice # noqa from chalice.app import CORSConfig # noqa @@ -47,7 +56,10 @@ def time(self): def create_local_server(app_obj, config, host, port): # type: (Chalice, Config, str, int) -> LocalDevServer - return LocalDevServer(app_obj, config, host, port) + local_app_opj = cast(LocalChalice, app_obj) # type: LocalChalice + local_app_opj.__class__ = LocalChalice + local_app_opj.initialize_current_request_lookup() + return LocalDevServer(local_app_opj, config, host, port) class LocalARNBuilder(object): @@ -661,3 +673,28 @@ def shutdown(self): # type: () -> None if self._server is not None: self._server.shutdown() + + +class LocalChalice(Chalice): + def __init__(self, *args, **kwargs): + # type: (Any, Any) -> None + self._current_request_lookup = {} # type: Dict[int, Optional[Request]] + super(LocalChalice, self).__init__(*args, **kwargs) + + def initialize_current_request_lookup(self): + # type: () -> None + self._current_request_lookup = {} + + @property + def current_request(self): # noqa + # type: () -> Optional[Request] + thread_id = threading.current_thread().ident + assert thread_id is not None + return self._current_request_lookup.get(thread_id, None) + + @current_request.setter + def current_request(self, value): # noqa + # type: (Optional[Request]) -> None + thread_id = threading.current_thread().ident + assert thread_id is not None + self._current_request_lookup[thread_id] = value