Coverage for src/httpx/_streams.py: 96%
76 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-06-09 10:38 +0100
« prev ^ index » next coverage.py v7.6.12, created at 2025-06-09 10:38 +0100
1import os
2import typing
5__all__ = ["ByteStream", "FileStream", "IterByteStream", "Stream"]
8def humanize_size(b: int) -> str:
9 B = float(b)
10 KB = float(1024)
11 MB = float(KB**2) # 1,048,576
12 GB = float(KB**3) # 1,073,741,824
13 TB = float(KB**4) # 1,099,511,627,776
15 if B < KB:
16 return "{0:.0f}B".format(B)
17 elif KB <= B < MB:
18 return "{0:.0f}KB".format(B / KB)
19 elif MB <= B < GB:
20 return "{0:.0f}MB".format(B / MB)
21 elif GB <= B < TB:
22 return "{0:.0f}GB".format(B / GB)
23 return "{0:.0f}TB".format(B / TB)
26class Stream:
27 @property
28 def size(self) -> int | None:
29 raise NotImplementedError()
31 def __iter__(self):
32 raise NotImplementedError()
35class ByteStream(Stream):
36 def __init__(self, content: bytes) -> None:
37 self._content = content
39 def __eq__(self, other) -> bool:
40 return (
41 self._content == other._content
42 if isinstance(other, ByteStream)
43 else False
44 )
46 @property
47 def size(self):
48 return len(self._content)
50 def __iter__(self) -> typing.Iterator[bytes]:
51 yield self._content
53 def __repr__(self) -> str:
54 size = humanize_size(self.size)
55 return f"<ByteStream [{size}]>"
58class IterByteStream(Stream):
59 def __init__(
60 self,
61 iterator: typing.Iterator[bytes],
62 size: int | None = None
63 ) -> None:
64 self._iterator = iterator
65 self._size = size
66 self._consumed = 0
68 @property
69 def size(self) -> int | None:
70 return self._size
72 def __iter__(self) -> typing.Iterator[bytes]:
73 self._consumed = 0
74 for part in self._iterator:
75 self._consumed += len(part)
76 if self._size is not None and self._consumed > self._size:
77 raise ValueError("IterByteStream returned more data than expected size")
78 yield part
80 if self._size is not None and self._consumed < self._size:
81 raise ValueError("IterByteStream returned less data than expected size")
82 self._size = self._consumed
84 def __repr__(self) -> str:
85 if self._size is None:
86 percent_of_size = "0% of ???" if self._consumed == 0 else "???% of ???"
87 else:
88 percent = self._consumed * 100 // self._size
89 size = humanize_size(self._size)
90 percent_of_size = f"{percent}% of {size}"
92 return f"<IterByteStream [{percent_of_size}]>"
95class FileStream(Stream):
96 def __init__(self, path: str):
97 self._path = path
99 @property
100 def size(self) -> int | None:
101 return os.path.getsize(self._path)
103 def __iter__(self) -> typing.Iterator[bytes]:
104 BUFFER_SIZE = 64 * 1024
105 with open(self._path, "rb") as fin:
106 while buffer := fin.read(BUFFER_SIZE):
107 yield buffer
109 def __repr__(self):
110 size = humanize_size(self.size)
111 return f"<FileStream [{self._path!r} {size}]>"