diff --git a/jupyter_server_proxy/handlers.py b/jupyter_server_proxy/handlers.py index 0a88eb39..89d4b8b4 100644 --- a/jupyter_server_proxy/handlers.py +++ b/jupyter_server_proxy/handlers.py @@ -4,7 +4,8 @@ Some original inspiration from https://github.com/senko/tornado-proxy """ -import os, json, re +import os +import re import socket from asyncio import Lock from copy import copy @@ -305,7 +306,7 @@ def _build_proxy_request(self, host, port, proxied_path, body, **extra_opts): decompress_response=False, headers=headers, **self.proxy_request_options(), - **extra_opts, + **extra_opts, ) return req @@ -374,12 +375,12 @@ async def proxy(self, host, port, proxied_path): else: client = httpclient.AsyncHTTPClient(force_instance=True) # check if the request is stream request - accept_header = self.request.headers.get('Accept') - if accept_header == 'text/event-stream': + accept_header = self.request.headers.get("Accept") + if accept_header == "text/event-stream": return await self._proxy_progressive(host, port, proxied_path, body, client) else: return await self._proxy_buffered(host, port, proxied_path, body, client) - + async def _proxy_progressive(self, host, port, proxied_path, body, client): # Proxy in progressive flush mode, whenever chunks are received. Potentially slower but get results quicker for voila # Set up handlers so we can progressively flush result @@ -388,15 +389,19 @@ async def _proxy_progressive(self, host, port, proxied_path, body, client): def dump_headers(headers_raw): for line in headers_raw: - r = re.match('^([a-zA-Z0-9\-_]+)\s*\:\s*([^\r\n]+)[\r\n]*$', line) + r = re.match("^([a-zA-Z0-9\\-_]+)\\s*\\:\\s*([^\r\n]+)[\r\n]*$", line) if r: - k,v = r.groups([1,2]) - if k not in ('Content-Length', 'Transfer-Encoding', - 'Content-Encoding', 'Connection'): + k, v = r.groups([1, 2]) + if k not in ( + "Content-Length", + "Transfer-Encoding", + "Content-Encoding", + "Connection", + ): # some header appear multiple times, eg 'Set-Cookie' - self.set_header(k,v) + self.set_header(k, v) else: - r = re.match('^HTTP[^\s]* ([0-9]+)', line) + r = re.match(r"^HTTP[^\s]* ([0-9]+)", line) if r: status_code = r.group(1) self.set_status(int(status_code)) @@ -412,20 +417,27 @@ def streaming_callback(chunk): # record activity at start and end of requests self._record_activity() # Do this here, not in header_callback so we can be sure headers are out of the way first - dump_headers(headers_raw) # array will be empty if this was already called before + dump_headers( + headers_raw + ) # array will be empty if this was already called before self.write(chunk) self.flush() # Now make the request - req = self._build_proxy_request(host, port, proxied_path, body, - streaming_callback=streaming_callback, - header_callback=header_callback) - + req = self._build_proxy_request( + host, + port, + proxied_path, + body, + streaming_callback=streaming_callback, + header_callback=header_callback, + ) + # no timeout for stream api req.request_timeout = 7200 req.connect_timeout = 600 - + try: response = await client.fetch(req, raise_error=False) except httpclient.HTTPError as err: @@ -442,15 +454,16 @@ def streaming_callback(chunk): self.set_status(500) self.write(str(response.error)) else: - self.set_status(response.code, response.reason) # Should already have been set + self.set_status( + response.code, response.reason + ) # Should already have been set - dump_headers(headers_raw) # Should already have been emptied + dump_headers(headers_raw) # Should already have been emptied - if response.body: # Likewise, should already be chunked out and flushed + if response.body: # Likewise, should already be chunked out and flushed self.write(response.body) async def _proxy_buffered(self, host, port, proxied_path, body, client): - req = self._build_proxy_request(host, port, proxied_path, body) self.log.debug(f"Proxying request to {req.url}") @@ -533,7 +546,6 @@ def rewrite_pe(rewritable_response: RewritableResponse): if rewritten_response.body: self.write(rewritten_response.body) - async def proxy_open(self, host, port, proxied_path=""): """ Called when a client opens a websocket connection. diff --git a/tests/resources/jupyter_server_config.py b/tests/resources/jupyter_server_config.py index c3cb296e..82c35558 100644 --- a/tests/resources/jupyter_server_config.py +++ b/tests/resources/jupyter_server_config.py @@ -42,10 +42,10 @@ def cats_only(response, path): response.code = 403 response.body = b"dogs not allowed" + def my_env(): - return { - "MYVAR": "String with escaped {{var}}" - } + return {"MYVAR": "String with escaped {{var}}"} + c.ServerProxy.servers = { "python-http": { diff --git a/tests/test_proxies.py b/tests/test_proxies.py index 515a081b..f0b59ecc 100644 --- a/tests/test_proxies.py +++ b/tests/test_proxies.py @@ -375,7 +375,7 @@ def streaming_cb(data): ) assert times_called == limit assert all([0.45 < t < 3.0 for t in stream_read_intervals]) - assert stream_data == [b'data: 0\n\n', b'data: 1\n\n', b'data: 2\n\n'] + assert stream_data == [b"data: 0\n\n", b"data: 1\n\n", b"data: 2\n\n"] async def test_server_proxy_websocket_messages(