Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain/plain/internal/files/uploadhandler.py: 55%

82 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-23 11:16 -0600

1""" 

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

3""" 

4 

5import os 

6from io import BytesIO 

7 

8from plain.internal.files.uploadedfile import ( 

9 InMemoryUploadedFile, 

10 TemporaryUploadedFile, 

11) 

12from plain.runtime import settings 

13from plain.utils.module_loading import import_string 

14 

15__all__ = [ 

16 "UploadFileException", 

17 "StopUpload", 

18 "SkipFile", 

19 "FileUploadHandler", 

20 "TemporaryFileUploadHandler", 

21 "MemoryFileUploadHandler", 

22 "load_handler", 

23 "StopFutureHandlers", 

24] 

25 

26 

27class UploadFileException(Exception): 

28 """ 

29 Any error having to do with uploading files. 

30 """ 

31 

32 pass 

33 

34 

35class StopUpload(UploadFileException): 

36 """ 

37 This exception is raised when an upload must abort. 

38 """ 

39 

40 def __init__(self, connection_reset=False): 

41 """ 

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

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

44 show a "connection reset" error. 

45 """ 

46 self.connection_reset = connection_reset 

47 

48 def __str__(self): 

49 if self.connection_reset: 

50 return "StopUpload: Halt current upload." 

51 else: 

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

53 

54 

55class SkipFile(UploadFileException): 

56 """ 

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

58 """ 

59 

60 pass 

61 

62 

63class StopFutureHandlers(UploadFileException): 

64 """ 

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

66 run should raise this exception instead of returning None. 

67 """ 

68 

69 pass 

70 

71 

72class FileUploadHandler: 

73 """ 

74 Base class for streaming upload handlers. 

75 """ 

76 

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

78 

79 def __init__(self, request=None): 

80 self.file_name = None 

81 self.content_type = None 

82 self.content_length = None 

83 self.charset = None 

84 self.content_type_extra = None 

85 self.request = request 

86 

87 def handle_raw_input( 

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

89 ): 

90 """ 

91 Handle the raw input from the client. 

92 

93 Parameters: 

94 

95 :input_data: 

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

97 :META: 

98 ``request.META``. 

99 :content_length: 

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

101 client. 

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

103 prepend two '--'. 

104 """ 

105 pass 

106 

107 def new_file( 

108 self, 

109 field_name, 

110 file_name, 

111 content_type, 

112 content_length, 

113 charset=None, 

114 content_type_extra=None, 

115 ): 

116 """ 

117 Signal that a new file has been started. 

118 

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

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

121 """ 

122 self.field_name = field_name 

123 self.file_name = file_name 

124 self.content_type = content_type 

125 self.content_length = content_length 

126 self.charset = charset 

127 self.content_type_extra = content_type_extra 

128 

129 def receive_data_chunk(self, raw_data, start): 

130 """ 

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

132 in the file of the chunk. 

133 """ 

134 raise NotImplementedError( 

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

136 ) 

137 

138 def file_complete(self, file_size): 

139 """ 

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

141 size accumulated by all the chunks. 

142 

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

144 """ 

145 raise NotImplementedError( 

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

147 ) 

148 

149 def upload_complete(self): 

150 """ 

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

152 that is necessary for this handler. 

153 """ 

154 pass 

155 

156 def upload_interrupted(self): 

157 """ 

158 Signal that the upload was interrupted. Subclasses should perform 

159 cleanup that is necessary for this handler. 

160 """ 

161 pass 

162 

163 

164class TemporaryFileUploadHandler(FileUploadHandler): 

165 """ 

166 Upload handler that streams data into a temporary file. 

167 """ 

168 

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

170 """ 

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

172 """ 

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

174 self.file = TemporaryUploadedFile( 

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

176 ) 

177 

178 def receive_data_chunk(self, raw_data, start): 

179 self.file.write(raw_data) 

180 

181 def file_complete(self, file_size): 

182 self.file.seek(0) 

183 self.file.size = file_size 

184 return self.file 

185 

186 def upload_interrupted(self): 

187 if hasattr(self, "file"): 

188 temp_location = self.file.temporary_file_path() 

189 try: 

190 self.file.close() 

191 os.remove(temp_location) 

192 except FileNotFoundError: 

193 pass 

194 

195 

196class MemoryFileUploadHandler(FileUploadHandler): 

197 """ 

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

199 """ 

200 

201 def handle_raw_input( 

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

203 ): 

204 """ 

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

206 used. 

207 """ 

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

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

210 self.activated = content_length <= settings.FILE_UPLOAD_MAX_MEMORY_SIZE 

211 

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

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

214 if self.activated: 

215 self.file = BytesIO() 

216 raise StopFutureHandlers() 

217 

218 def receive_data_chunk(self, raw_data, start): 

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

220 if self.activated: 

221 self.file.write(raw_data) 

222 else: 

223 return raw_data 

224 

225 def file_complete(self, file_size): 

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

227 if not self.activated: 

228 return 

229 

230 self.file.seek(0) 

231 return InMemoryUploadedFile( 

232 file=self.file, 

233 field_name=self.field_name, 

234 name=self.file_name, 

235 content_type=self.content_type, 

236 size=file_size, 

237 charset=self.charset, 

238 content_type_extra=self.content_type_extra, 

239 ) 

240 

241 

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

243 """ 

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

245 

246 E.g.:: 

247 >>> from plain.http import HttpRequest 

248 >>> request = HttpRequest() 

249 >>> load_handler( 

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

251 ... request, 

252 ... ) 

253 <TemporaryFileUploadHandler object at 0x...> 

254 """ 

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