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

1import os 

2import typing 

3 

4 

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

6 

7 

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 

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 

26class Stream: 

27 @property 

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

29 raise NotImplementedError() 

30 

31 def __iter__(self): 

32 raise NotImplementedError() 

33 

34 

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

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 

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 

95class 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}]>"