code style improvements

This commit is contained in:
Maximilian Hils 2021-03-15 17:23:45 +01:00
parent 12e4785d44
commit 7fae3c8951
4 changed files with 12 additions and 27 deletions

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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)