typing++

Author: Kriechi

src/hpack/hpack.py

@@ -5,7 +5,7 @@
 Implements the HPACK header compression algorithm as detailed by the IETF.
 """
 import logging
-from typing import Any, Generator, Union
+from typing import Any, Generator, Iterable, Optional, Union
 
 from .table import HeaderTable, table_entry_size
 from .exceptions import (
@@ -16,7 +16,8 @@
     REQUEST_CODES, REQUEST_CODES_LENGTH
 )
 from .huffman_table import decode_huffman
-from .struct import HeaderTuple, NeverIndexedHeaderTuple, Headers
+from .struct import HeaderTuple, NeverIndexedHeaderTuple, HeaderWeaklyTyped
+from .table import HeaderTable, table_entry_size
 
 log = logging.getLogger(__name__)
 
@@ -34,17 +35,17 @@
 DEFAULT_MAX_HEADER_LIST_SIZE = 2 ** 16
 
 
-def _unicode_if_needed(header: HeaderTuple, raw: bool) -> HeaderTuple:
+def _unicode_if_needed(header: HeaderWeaklyTyped, raw: bool) -> HeaderTuple:
     """
     Provides a header as a unicode string if raw is False, otherwise returns
     it as a bytestring.
     """
     name = bytes(header[0])  # type: ignore
     value = bytes(header[1])  # type: ignore
+
     if not raw:
-        return header.__class__(name.decode('utf-8'), value.decode('utf-8'))
-    else:
-        return header.__class__(name, value)
+        return header.__class__(name.decode("utf-8"), value.decode("utf-8"))  # type: ignore
+    return header.__class__(name, value)  # type: ignore
 
 
 def encode_integer(integer: int, prefix_bits: int) -> bytearray:
@@ -123,7 +124,7 @@ def decode_integer(data: bytes, prefix_bits: int) -> tuple[int, int]:
 
 
 def _dict_to_iterable(header_dict: Union[dict[bytes, bytes], dict[str, str]]) \
-        -> Generator[Union[tuple[bytes, bytes], tuple[str, str]], None, None]:
+        -> Generator[Union[tuple[bytes, bytes, Optional[bool]], tuple[str, str, Optional[bool]]], None, None]:
     """
     This converts a dictionary to an iterable of two-tuples. This is a
     HPACK-specific function because it pulls "special-headers" out first and
@@ -177,7 +178,7 @@ def header_table_size(self, value: int) -> None:
             self.table_size_changes.append(value)
 
     def encode(self,
-               headers: Headers,
+               headers: Union[Iterable[tuple[Union[bytes, str], Union[bytes, str], Optional[bool]]], dict[Union[bytes, str], Union[bytes, str, tuple]]],
                huffman: bool = True) -> bytes:
         """
         Takes a set of headers and encodes them into a HPACK-encoded header
@@ -433,7 +434,7 @@ def header_table_size(self) -> int:
     def header_table_size(self, value: int) -> None:
         self.header_table.maxsize = value
 
-    def decode(self, data: bytes, raw: bool = False) -> Headers:
+    def decode(self, data: bytes, raw: bool = False) -> Iterable[HeaderTuple]:
         """
         Takes an HPACK-encoded header block and decodes it into a header set.
 

src/hpack/huffman.py

@@ -7,6 +7,9 @@
 """
 
 
+from typing import Optional
+
+
 class HuffmanEncoder:
     """
     Encodes a string according to the Huffman encoding table defined in the
@@ -16,7 +19,7 @@ def __init__(self, huffman_code_list: list[int], huffman_code_list_lengths: list
         self.huffman_code_list = huffman_code_list
         self.huffman_code_list_lengths = huffman_code_list_lengths
 
-    def encode(self, bytes_to_encode: bytes) -> bytes:
+    def encode(self, bytes_to_encode: Optional[bytes]) -> bytes:
         """
         Given a string of bytes, encodes them according to the HPACK Huffman
         specification.

src/hpack/struct.py

@@ -5,11 +5,12 @@
 Contains structures for representing header fields with associated metadata.
 """
 
+from typing import Any
 
-from typing import Any, Iterable, Union
+from typing_extensions import Self, TypeAlias
 
 
-class HeaderTuple(tuple[Union[bytes, str], Union[bytes, str]]):
+class HeaderTuple(tuple[bytes, bytes]):
     """
     A data structure that stores a single header field.
 
@@ -27,7 +28,7 @@ class HeaderTuple(tuple[Union[bytes, str], Union[bytes, str]]):
 
     indexable = True
 
-    def __new__(cls, *args: Any) -> "HeaderTuple":
+    def __new__(cls, *args: Any) -> Self:
         return tuple.__new__(cls, args)
 
 
@@ -40,10 +41,9 @@ class NeverIndexedHeaderTuple(HeaderTuple):
 
     indexable = False
 
-    def __new__(cls, *args: Any) -> "NeverIndexedHeaderTuple":
+    def __new__(cls, *args: Any) -> Self:
         return tuple.__new__(cls, args)
 
 
-# explicitly mentioning all valid header types , even if some superseed each other
-Header = Union[HeaderTuple, NeverIndexedHeaderTuple, tuple[bytes, bytes], tuple[str, str]]
-Headers = Iterable[Header]
+Header: TypeAlias = "HeaderTuple | NeverIndexedHeaderTuple | tuple[bytes, bytes]"
+HeaderWeaklyTyped: TypeAlias = "HeaderTuple | NeverIndexedHeaderTuple | tuple[bytes | str, bytes | str]"

src/hpack/table.py

@@ -1,14 +1,14 @@
 # flake8: noqa
 from collections import deque
 import logging
-from typing import Any, Optional, Union
+from typing import Optional
 
 from .exceptions import InvalidTableIndex
 
 log = logging.getLogger(__name__)
 
 
-def table_entry_size(name: Union[bytes, str], value: Union[bytes, str]) -> int:
+def table_entry_size(name: bytes, value: bytes) -> int:
     """
     Calculates the size of a single entry