From 7fae3c895192ca65be69390cd44e5ca51fcfc6b4 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 15 Mar 2021 17:23:45 +0100 Subject: [PATCH] code style improvements --- mitmproxy/proxy/layers/http/__init__.py | 1 - mitmproxy/proxy/layers/http/_http2.py | 20 +++----------------- test/mitmproxy/proxy/conftest.py | 2 +- test/mitmproxy/proxy/tutils.py | 16 ++++++++-------- 4 files changed, 12 insertions(+), 27 deletions(-) diff --git a/mitmproxy/proxy/layers/http/__init__.py b/mitmproxy/proxy/layers/http/__init__.py index eedb0ec65..dd846f847 100644 --- a/mitmproxy/proxy/layers/http/__init__.py +++ b/mitmproxy/proxy/layers/http/__init__.py @@ -5,7 +5,6 @@ from dataclasses import dataclass from typing import DefaultDict, Dict, List, Optional, Tuple, Union import wsproto.handshake - from mitmproxy import flow, http from mitmproxy.connection import Connection, Server from mitmproxy.net import server_spec diff --git a/mitmproxy/proxy/layers/http/_http2.py b/mitmproxy/proxy/layers/http/_http2.py index 73d7b8d0f..f2b5d5a12 100644 --- a/mitmproxy/proxy/layers/http/_http2.py +++ b/mitmproxy/proxy/layers/http/_http2.py @@ -45,10 +45,6 @@ class Http2Connection(HttpConnection): streams: Dict[int, StreamState] """keep track of all active stream ids to send protocol errors on teardown""" - SendProtocolError: Type[Union[RequestProtocolError, ResponseProtocolError]] - SendData: Type[Union[RequestData, ResponseData]] - SendEndOfMessage: Type[Union[RequestEndOfMessage, ResponseEndOfMessage]] - ReceiveProtocolError: Type[Union[RequestProtocolError, ResponseProtocolError]] ReceiveData: Type[Union[RequestData, ResponseData]] ReceiveTrailers: Type[Union[RequestTrailers, ResponseTrailers]] @@ -94,19 +90,17 @@ class Http2Connection(HttpConnection): yield SendData(self.conn, self.h2_conn.data_to_send()) elif isinstance(event, HttpEvent): - if isinstance(event, self.SendData): - assert isinstance(event, (RequestData, ResponseData)) + if isinstance(event, (RequestData, ResponseData)): if self.is_open_for_us(event.stream_id): self.h2_conn.send_data(event.stream_id, event.data) elif isinstance(event, (RequestTrailers, ResponseTrailers)): if self.is_open_for_us(event.stream_id): trailers = [*event.trailers.fields] self.h2_conn.send_headers(event.stream_id, trailers, end_stream=True) - elif isinstance(event, self.SendEndOfMessage): + elif isinstance(event, (RequestEndOfMessage, ResponseEndOfMessage)): if self.is_open_for_us(event.stream_id): self.h2_conn.end_stream(event.stream_id) - elif isinstance(event, self.SendProtocolError): - assert isinstance(event, (RequestProtocolError, ResponseProtocolError)) + elif isinstance(event, (RequestProtocolError, ResponseProtocolError)): if not self.is_closed(event.stream_id): code = { status_codes.CLIENT_CLOSED_REQUEST: h2.errors.ErrorCodes.CANCEL, @@ -274,10 +268,6 @@ class Http2Server(Http2Connection): client_side=False, ) - SendProtocolError = ResponseProtocolError - SendData = ResponseData - SendEndOfMessage = ResponseEndOfMessage - ReceiveProtocolError = RequestProtocolError ReceiveData = RequestData ReceiveTrailers = RequestTrailers @@ -339,10 +329,6 @@ class Http2Client(Http2Connection): client_side=True, ) - SendProtocolError = RequestProtocolError - SendData = RequestData - SendEndOfMessage = RequestEndOfMessage - ReceiveProtocolError = ResponseProtocolError ReceiveData = ResponseData ReceiveTrailers = ResponseTrailers diff --git a/test/mitmproxy/proxy/conftest.py b/test/mitmproxy/proxy/conftest.py index 6dde89ffa..8477158f3 100644 --- a/test/mitmproxy/proxy/conftest.py +++ b/test/mitmproxy/proxy/conftest.py @@ -3,7 +3,7 @@ import os import pytest from hypothesis import settings -from mitmproxy import options, connection +from mitmproxy import connection, options from mitmproxy.addons.core import Core from mitmproxy.addons.proxyserver import Proxyserver from mitmproxy.addons.termlog import TermLog diff --git a/test/mitmproxy/proxy/tutils.py b/test/mitmproxy/proxy/tutils.py index 0a38985af..aad63d8ca 100644 --- a/test/mitmproxy/proxy/tutils.py +++ b/test/mitmproxy/proxy/tutils.py @@ -318,7 +318,10 @@ class reply(events.Event): return inst -class _Placeholder: +T = typing.TypeVar("T") + + +class _Placeholder(typing.Generic[T]): """ Placeholder value in playbooks, so that objects (flows in particular) can be referenced before they are known. Example: @@ -333,15 +336,15 @@ class _Placeholder: assert f().messages == 0 """ - def __init__(self, cls: typing.Type): + def __init__(self, cls: typing.Type[T]): self._obj = None self._cls = cls - def __call__(self): + def __call__(self) -> T: """Get the actual object""" return self._obj - def setdefault(self, value): + def setdefault(self, value: T) -> T: if self._obj is None: if self._cls is not typing.Any and not isinstance(value, self._cls): raise TypeError(f"expected {self._cls.__name__}, got {type(value).__name__}.") @@ -355,11 +358,8 @@ class _Placeholder: return f"Placeholder:{str(self._obj)}" -T = typing.TypeVar("T") - - # noinspection PyPep8Naming -def Placeholder(cls: typing.Type[T] = typing.Any) -> typing.Union[T, typing.Callable[[], T]]: +def Placeholder(cls: typing.Type[T] = typing.Any) -> typing.Union[T, _Placeholder[T]]: return _Placeholder(cls)