Coverage for src/httpx/_streams.py: 96%

1 import os

2 import typing

3

4

5 __all__ = ["ByteStream", "FileStream", "IterByteStream", "Stream"]

6

7

8 def 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

14

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)

24

25

26 class Stream:

27 @property

28 def size(self) -> int | None:

29 raise NotImplementedError()

30

31 def __iter__(self):

32 raise NotImplementedError()

33

34

35 class ByteStream(Stream):

36 def __init__(self, content: bytes) -> None:

37 self._content = content

38

39 def __eq__(self, other) -> bool:

40 return (

41 self._content == other._content

42 if isinstance(other, ByteStream)

43 else False

44 )

45

46 @property

47 def size(self):

48 return len(self._content)

49

50 def __iter__(self) -> typing.Iterator[bytes]:

51 yield self._content

52

53 def __repr__(self) -> str:

54 size = humanize_size(self.size)

55 return f"<ByteStream [{size}]>"

56

57

58 class 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

67

68 @property

69 def size(self) -> int | None:

70 return self._size

71

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

79

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

83

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}"

91

92 return f"<IterByteStream [{percent_of_size}]>"

93

94

95 class FileStream(Stream):

96 def __init__(self, path: str):

97 self._path = path

98

99 @property

100 def size(self) -> int | None:

101 return os.path.getsize(self._path)

102

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

108

109 def __repr__(self):

110 size = humanize_size(self.size)

111 return f"<FileStream [{self._path! r} {size}]>"