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

1"""

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

3"""

4

5 import os

6 from io import BytesIO

7

8 from plain.internal.files.uploadedfile import (

9 InMemoryUploadedFile,

10 TemporaryUploadedFile,

11)

12 from plain.runtime import settings

13 from 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

27 class UploadFileException(Exception):

28 """

29 Any error having to do with uploading files.

30 """

31

32 pass

33

34

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

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

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

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

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

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

242 def 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)