Test suite runs and passes!
[zzz-floof.git] / floof / lib / file_storage.py
1 # This sucks. Make it not suck.
2 import os, shutil, tempfile, hashlib
3
4 chunk_size = 1024*1024 # TODO: is this a good chunk size?
5
6 from pylons import config
7
8 """
9 Notes:
10 # You can get a mime type like so:
11 from mimetypes import guess_type
12 guess_type(temp.filename)[0]
13 """
14
15 def get_path(space, hash):
16 return "/" + os.path.join( space, hash[:2], hash[2:] )
17
18
19 def save_file(space, fileobj, hash=None):
20 """Saves the contents of fileobj to the given storage space.
21
22 If a hash is not provided, the SHA-1 sum of the file will be computed and
23 that will be used. The ideal scenario here is to let the hash of the
24 original file be computed automatically, then create thumbnails et al. with
25 the same hash in a different space.
26
27 Returns the hashsum.
28 """
29 dest_root = os.path.join( config['app_conf']['static_root'], space )
30
31 # The incoming fileobj could be a tempfile that's already been unlinked --
32 # and probably is, as it's coming from a Pylons upload object. Thus we
33 # have to copy the data rather than the file, and we may have to read the
34 # file anyway to hash it, so do both at the same time.
35
36 # Need a named temporary file to write to; it gets renamed once the hash is
37 # computed
38 temp_fd, temp_path = tempfile.mkstemp()
39 temp_file = os.fdopen(temp_fd, "wb")
40
41 sha1 = hashlib.sha1()
42 while True:
43 data = fileobj.read(chunk_size)
44 if not data:
45 break
46
47 if not hash:
48 sha1.update(data)
49 temp_file.write(data)
50
51 temp_file.close()
52 if not hash:
53 hash = sha1.hexdigest()
54
55 # Git convention: first two characters are the directory
56 dest_dir = os.path.join(dest_root, hash[:2])
57 dest_path = os.path.join(dest_dir, hash[2:])
58
59 if not os.path.exists(dest_dir):
60 os.makedirs(dest_dir)
61 print dest_path
62 shutil.move(temp_path, dest_path)
63
64 return hash