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

1"""

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

3"""

4 import os

5 from io import BytesIO

6

7 from plain.internal.files.uploadedfile import (

8 InMemoryUploadedFile,

9 TemporaryUploadedFile,

10)

11 from plain.runtime import settings

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

26 class UploadFileException(Exception):

27 """

28 Any error having to do with uploading files.

29 """

30

31 pass

32

33

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

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

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

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

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

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

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