Coverage for /Users/davegaeddert/Development/dropseed/plain/plain/plain/internal/files/uploadhandler.py: 43%

82 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-16 22:04 -0500

1""" 

2Base file upload handler classes, and the built-in concrete subclasses 

3""" 

4import os 

5from io import BytesIO 

6 

7from plain.internal.files.uploadedfile import ( 

8 InMemoryUploadedFile, 

9 TemporaryUploadedFile, 

10) 

11from plain.runtime import settings 

12from plain.utils.module_loading import import_string 

13 

14__all__ = [ 

15 "UploadFileException", 

16 "StopUpload", 

17 "SkipFile", 

18 "FileUploadHandler", 

19 "TemporaryFileUploadHandler", 

20 "MemoryFileUploadHandler", 

21 "load_handler", 

22 "StopFutureHandlers", 

23] 

24 

25 

26class UploadFileException(Exception): 

27 """ 

28 Any error having to do with uploading files. 

29 """ 

30 

31 pass 

32 

33 

34class StopUpload(UploadFileException): 

35 """ 

36 This exception is raised when an upload must abort. 

37 """ 

38 

39 def __init__(self, connection_reset=False): 

40 """ 

41 If ``connection_reset`` is ``True``, Plain knows will halt the upload 

42 without consuming the rest of the upload. This will cause the browser to 

43 show a "connection reset" error. 

44 """ 

45 self.connection_reset = connection_reset 

46 

47 def __str__(self): 

48 if self.connection_reset: 

49 return "StopUpload: Halt current upload." 

50 else: 

51 return "StopUpload: Consume request data, then halt." 

52 

53 

54class SkipFile(UploadFileException): 

55 """ 

56 This exception is raised by an upload handler that wants to skip a given file. 

57 """ 

58 

59 pass 

60 

61 

62class StopFutureHandlers(UploadFileException): 

63 """ 

64 Upload handlers that have handled a file and do not want future handlers to 

65 run should raise this exception instead of returning None. 

66 """ 

67 

68 pass 

69 

70 

71class FileUploadHandler: 

72 """ 

73 Base class for streaming upload handlers. 

74 """ 

75 

76 chunk_size = 64 * 2**10 # : The default chunk size is 64 KB. 

77 

78 def __init__(self, request=None): 

79 self.file_name = None 

80 self.content_type = None 

81 self.content_length = None 

82 self.charset = None 

83 self.content_type_extra = None 

84 self.request = request 

85 

86 def handle_raw_input( 

87 self, input_data, META, content_length, boundary, encoding=None 

88 ): 

89 """ 

90 Handle the raw input from the client. 

91 

92 Parameters: 

93 

94 :input_data: 

95 An object that supports reading via .read(). 

96 :META: 

97 ``request.META``. 

98 :content_length: 

99 The (integer) value of the Content-Length header from the 

100 client. 

101 :boundary: The boundary from the Content-Type header. Be sure to 

102 prepend two '--'. 

103 """ 

104 pass 

105 

106 def new_file( 

107 self, 

108 field_name, 

109 file_name, 

110 content_type, 

111 content_length, 

112 charset=None, 

113 content_type_extra=None, 

114 ): 

115 """ 

116 Signal that a new file has been started. 

117 

118 Warning: As with any data from the client, you should not trust 

119 content_length (and sometimes won't even get it). 

120 """ 

121 self.field_name = field_name 

122 self.file_name = file_name 

123 self.content_type = content_type 

124 self.content_length = content_length 

125 self.charset = charset 

126 self.content_type_extra = content_type_extra 

127 

128 def receive_data_chunk(self, raw_data, start): 

129 """ 

130 Receive data from the streamed upload parser. ``start`` is the position 

131 in the file of the chunk. 

132 """ 

133 raise NotImplementedError( 

134 "subclasses of FileUploadHandler must provide a receive_data_chunk() method" 

135 ) 

136 

137 def file_complete(self, file_size): 

138 """ 

139 Signal that a file has completed. File size corresponds to the actual 

140 size accumulated by all the chunks. 

141 

142 Subclasses should return a valid ``UploadedFile`` object. 

143 """ 

144 raise NotImplementedError( 

145 "subclasses of FileUploadHandler must provide a file_complete() method" 

146 ) 

147 

148 def upload_complete(self): 

149 """ 

150 Signal that the upload is complete. Subclasses should perform cleanup 

151 that is necessary for this handler. 

152 """ 

153 pass 

154 

155 def upload_interrupted(self): 

156 """ 

157 Signal that the upload was interrupted. Subclasses should perform 

158 cleanup that is necessary for this handler. 

159 """ 

160 pass 

161 

162 

163class TemporaryFileUploadHandler(FileUploadHandler): 

164 """ 

165 Upload handler that streams data into a temporary file. 

166 """ 

167 

168 def new_file(self, *args, **kwargs): 

169 """ 

170 Create the file object to append to as data is coming in. 

171 """ 

172 super().new_file(*args, **kwargs) 

173 self.file = TemporaryUploadedFile( 

174 self.file_name, self.content_type, 0, self.charset, self.content_type_extra 

175 ) 

176 

177 def receive_data_chunk(self, raw_data, start): 

178 self.file.write(raw_data) 

179 

180 def file_complete(self, file_size): 

181 self.file.seek(0) 

182 self.file.size = file_size 

183 return self.file 

184 

185 def upload_interrupted(self): 

186 if hasattr(self, "file"): 

187 temp_location = self.file.temporary_file_path() 

188 try: 

189 self.file.close() 

190 os.remove(temp_location) 

191 except FileNotFoundError: 

192 pass 

193 

194 

195class MemoryFileUploadHandler(FileUploadHandler): 

196 """ 

197 File upload handler to stream uploads into memory (used for small files). 

198 """ 

199 

200 def handle_raw_input( 

201 self, input_data, META, content_length, boundary, encoding=None 

202 ): 

203 """ 

204 Use the content_length to signal whether or not this handler should be 

205 used. 

206 """ 

207 # Check the content-length header to see if we should 

208 # If the post is too large, we cannot use the Memory handler. 

209 self.activated = content_length <= settings.FILE_UPLOAD_MAX_MEMORY_SIZE 

210 

211 def new_file(self, *args, **kwargs): 

212 super().new_file(*args, **kwargs) 

213 if self.activated: 

214 self.file = BytesIO() 

215 raise StopFutureHandlers() 

216 

217 def receive_data_chunk(self, raw_data, start): 

218 """Add the data to the BytesIO file.""" 

219 if self.activated: 

220 self.file.write(raw_data) 

221 else: 

222 return raw_data 

223 

224 def file_complete(self, file_size): 

225 """Return a file object if this handler is activated.""" 

226 if not self.activated: 

227 return 

228 

229 self.file.seek(0) 

230 return InMemoryUploadedFile( 

231 file=self.file, 

232 field_name=self.field_name, 

233 name=self.file_name, 

234 content_type=self.content_type, 

235 size=file_size, 

236 charset=self.charset, 

237 content_type_extra=self.content_type_extra, 

238 ) 

239 

240 

241def load_handler(path, *args, **kwargs): 

242 """ 

243 Given a path to a handler, return an instance of that handler. 

244 

245 E.g.:: 

246 >>> from plain.http import HttpRequest 

247 >>> request = HttpRequest() 

248 >>> load_handler( 

249 ... 'plain.internal.files.uploadhandler.TemporaryFileUploadHandler', 

250 ... request, 

251 ... ) 

252 <TemporaryFileUploadHandler object at 0x...> 

253 """ 

254 return import_string(path)(*args, **kwargs)