Add/correct missing/wrong type hints in stub files

fastwarc/fastwarc/stream_io.pyi

@@ -1,18 +1,26 @@
-from typing import ContextManager, IO
+from types import TracebackType
+from typing import ContextManager, IO, Optional, Type, Union, BinaryIO
 
 
-class IOStream(ContextManager):
+class IOStream(ContextManager[IOStream]):
     def read(self, size: int) -> bytes: ...
     def write(self, data: bytes) -> int: ...
     def close(self) -> None: ...
     def flush(self) -> None: ...
     def seek(self, offset: int) -> None: ...
     def tell(self) -> int: ...
+    def __enter__(self) -> IOStream: ...
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        traceback: Optional[TracebackType]
+    ) -> None: ...
 
 
 class BufferedReader:
     def __init__(
-        self, stream: IOStream, buf_size: int = 8192, negotiate_stream: bool = True
+        self, stream: Union[IOStream, BinaryIO], buf_size: int = 65536, negotiate_stream: bool = True
     ) -> None: ...
     def close(self) -> None: ...
     def consume(self, size: int = -1) -> int: ...
@@ -22,6 +30,7 @@ class BufferedReader:
 
 
 class BytesIOStream(IOStream):
+    def __init__(self, initial_data: Union[bytes, None] = None) -> None: ...
     def getvalue(self) -> bytes: ...
 
 
@@ -36,28 +45,28 @@ class CompressingStream(IOStream):
 
 class BrotliStream(CompressingStream):
     def __init__(
-        self, raw_stream: IOStream, quality: int = 11, lgwin: int = 22, lgblock: int = 0
+        self, raw_stream: Union[IOStream, BinaryIO], quality: int = 11, lgwin: int = 22, lgblock: int = 0
     ) -> None: ...
 
 
 class GZipStream(CompressingStream):
     def __init__(
-        self, raw_stream: IOStream, compression_level: int = 9, zlib: bool = False
+        self, raw_stream: Union[IOStream, BinaryIO], compression_level: int = 9, zlib: bool = False
     ) -> None: ...
 
 
 class LZ4Stream(CompressingStream):
     def __init__(
         self,
-        raw_stream: IOStream,
+        raw_stream: Union[IOStream, BinaryIO],
         compression_level: int = 12,
         favor_dec_speed: bool = True,
     ) -> None: ...
     def prepopulate(self, initial_data: bytes) -> None: ...
 
 
 class PythonIOStreamAdapter(IOStream):
-    def __init__(self, py_stream: IO) -> None: ...
+    def __init__(self, py_stream: BinaryIO) -> None: ...
 
 
 class FastWARCError(Exception):

fastwarc/fastwarc/warc.pyi

@@ -1,14 +1,16 @@
 from datetime import datetime
 from typing import (
+    Union,
     Optional,
     Iterator,
     Dict,
     Tuple,
-    MutableMapping,
+    Literal,
+    Callable,
     Iterable,
     ValuesView,
     KeysView,
-    Type,
+    BinaryIO,
 )
 from enum import IntFlag
 
@@ -33,20 +35,25 @@ no_type = WarcRecordType.no_type
 any_type = WarcRecordType.any_type
 
 
-class WarcHeaderMap(MutableMapping[str, str]):
+class WarcHeaderMap:
     reason_phrase: Optional[str]
     status_code: Optional[str]
     status_line: str
 
     def append(self, key: str, value: str) -> None: ...
     def asdict(self) -> Dict[str, str]: ...
-    def astuples(self) -> Tuple[str, str]: ...
+    def astuples(self) -> Tuple[Tuple[str, str], ...]: ...
     def clear(self) -> None: ...
     def get(self, key: str, default: Optional[str] = None) -> Optional[str]: ...
     def items(self) -> Iterator[Tuple[str, str]]: ...
     def keys(self) -> KeysView[str]: ...
     def values(self) -> ValuesView[str]: ...
-    def write(self, stream: IOStream) -> None: ...
+    def write(self, stream: Union[IOStream, BinaryIO]) -> None: ...
+    def __getitem__(self, item: str) -> str: ...
+    def __iter__(self) -> Iterator[Tuple[str, str]]: ...
+    def __len__(self) -> int: ...
+    def __setitem__(self, key: str, value: str) -> None: ...
+    def __contains__(self, item: str) -> bool: ...
 
 
 class WarcRecord:
@@ -59,32 +66,41 @@ class WarcRecord:
     is_http_parsed: bool
     http_headers: Optional[WarcHeaderMap]
     http_content_type: Optional[str]
-    http_content_type: Optional[str]
     http_charset: Optional[str]
     http_date: Optional[datetime]
     http_last_modified: Optional[datetime]
-    content_length: int
     reader: BufferedReader
     stream_pos: int
 
     def init_headers(
-        self, content_length: int = 0, record_type=no_type, record_urn=None
-    ): ...
+        self, content_length: int = 0, record_type: WarcRecordType = no_type, record_urn: Optional[bytes] = None
+    ) -> None: ...
     def freeze(self) -> bool: ...
     def set_bytes_content(self, content: bytes) -> None: ...
-    def parse_http(self, strict_mode=True, auto_decode: str = "none") -> None: ...
+    def parse_http(self, strict_mode: bool = True, auto_decode: str = "none") -> None: ...
     def verify_block_digest(self, consume: bool = False) -> bool: ...
     def verify_payload_digest(self, consume: bool = False) -> bool: ...
+    def write(
+        self,
+        stream: Union[IOStream, BinaryIO],
+        checksum_data: bool = False,
+        payload_digest: Optional[bytes] = None,
+        chunk_size: int = 16384
+    ) -> int: ...
 
 
 class ArchiveIterator(Iterable[WarcRecord]):
     def __init__(
         self,
-        stream: Type[IOStream],
+        stream: Union[IOStream, BinaryIO],
         record_types: WarcRecordType = any_type,
         parse_http: bool = True,
         min_content_length: int = -1,
         max_content_length: int = -1,
+        func_filter: Optional[Callable[[WarcRecord], bool]] = None,
+        verify_digests: bool = False,
+        strict_mode: bool = True,
+        auto_decode: Literal["none", "content", "transfer", "all"] = "none",
     ) -> None: ...
     def __iter__(self) -> Iterator[WarcRecord]: ...
     def __next__(self) -> WarcRecord: ...