Skip to content

Commit

Permalink
fix type hints
Browse files Browse the repository at this point in the history
  • Loading branch information
amyangfei committed Sep 12, 2022
1 parent 5c6f6d7 commit 03d60ff
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions gor/base.py
Expand Up @@ -7,6 +7,7 @@
import datetime
import traceback
from urllib.parse import quote_plus, urlparse, parse_qs
from typing import Union, Dict


def gor_hex_data(data):
Expand Down Expand Up @@ -114,7 +115,7 @@ def http_status(self, payload: bytes) -> str:
def set_http_status(self, payload: bytes, new_status: str) -> bytes:
return self.set_http_path(payload, new_status)

def http_headers(self, payload: bytes) -> dict[str, str]:
def http_headers(self, payload: Union[bytes, str]) -> Dict[str, str]:
"""
Parse the payload and return http headers in a map
:param payload: the http payload to inspect
Expand All @@ -134,7 +135,7 @@ def http_headers(self, payload: bytes) -> dict[str, str]:
headers[key] = value
return headers

def http_header(self, payload: bytes, name: str) -> dict[str, str]:
def http_header(self, payload: Union[bytes, str], name: str) -> Dict[str, str]:
current_line = 0
idx = 0
header = {
Expand Down Expand Up @@ -174,7 +175,7 @@ def http_header(self, payload: bytes, name: str) -> dict[str, str]:
idx += 1
return None

def set_http_header(self, payload: bytes, name: str, value: str) -> bytes:
def set_http_header(self, payload: Union[bytes, str], name: str, value: str) -> bytes:
if isinstance(payload, str):
payload = payload.encode()
header = self.http_header(payload, name)
Expand All @@ -184,7 +185,7 @@ def set_http_header(self, payload: bytes, name: str, value: str) -> bytes:
else:
return payload[:header['value_start']] + b' ' + value.encode() + b'\r\n' + payload[header['end']:]

def delete_http_header(self, payload: bytes, name: str) -> bytes:
def delete_http_header(self, payload: Union[bytes, str], name: str) -> bytes:
if isinstance(payload, str):
payload = payload.encode()
header = self.http_header(payload, name)
Expand Down

0 comments on commit 03d60ff

Please sign in to comment.