Fix stupid error if src is None.
[rox-archive.git] / formats.py
blob5812c4923fbb87ae250900f2eb4f7c455fb76a92
1 import os, sys
2 from support import PipeThroughCommand, escape, Tmp
3 import rox
5 current_command = None
7 def pipe_through_command(command, src, dst):
8 global current_command
9 assert not current_command
10 current_command = PipeThroughCommand(command, src, dst)
11 try:
12 current_command.wait()
13 finally:
14 current_command = None
16 operations = []
17 class Operation:
18 add_extension = 0
20 def __init__(self, extension):
21 operations.append(self)
22 self.extension = extension
24 def can_handle(self, data):
25 return isinstance(data, FileData)
27 def save_to_stream(self, data, stream):
28 pipe_through_command(self.command, data.source, stream)
30 class Compress(Operation):
31 "Compress a stream into another stream."
32 add_extension = 1
34 def __init__(self, extension, command, type):
35 Operation.__init__(self, extension)
36 self.command = command
37 self.type = type
39 def __str__(self):
40 return 'Compress as .%s' % self.extension
42 class Decompress(Operation):
43 "Decompress a stream into another stream."
44 type = 'text/plain'
46 def __init__(self, extension, command):
47 Operation.__init__(self, extension)
48 self.command = command
50 def __str__(self):
51 return 'Decompress .%s' % self.extension
53 class Extract(Operation):
54 "Extract an archive to a directory."
55 type = 'inode/directory'
57 def __init__(self, extension, command):
58 "If command has a %s then the source path is inserted, else uses stdin."
59 Operation.__init__(self, extension)
60 self.command = command
62 def __str__(self):
63 return 'Extract from a .%s' % self.extension
65 def save_to_stream(self, data, stream):
66 raise Exception('This operation creates a directory, so you have '
67 'to drag to a filer window on the local machine')
69 def save_to_file(self, data, path):
70 if os.path.exists(path):
71 if not os.path.isdir(path):
72 raise Exception("'%s' already exists and is not a directory!" %
73 path)
74 if not os.path.exists(path):
75 os.mkdir(path)
76 os.chdir(path)
77 command = self.command
78 source = data.source
79 if command.find("'%s'") != -1:
80 command = command % escape(source.name)
81 source = None
82 try:
83 pipe_through_command(command, source, None)
84 finally:
85 try:
86 os.rmdir(path) # Will only succeed if it's empty
87 except:
88 pass
89 # If we created only a single subdirectory, move it up.
90 dirs = os.listdir(path)
91 if len(dirs) != 1:
92 return
93 dir = dirs[0]
94 unneeded_path = os.path.join(path, dir)
95 if not os.path.isdir(unneeded_path):
96 return
97 import random
98 tmp_path = os.path.join(path, 'tmp-' + `random.randint(0, 100000)`)
99 os.rename(unneeded_path, tmp_path)
100 for file in os.listdir(tmp_path):
101 os.rename(os.path.join(tmp_path, file), os.path.join(path, file))
102 os.rmdir(tmp_path)
104 class Archive(Operation):
105 "Create an archive from a directory."
106 add_extension = 1
108 def __init__(self, extension, command, type):
109 assert command.find("'%s'") != -1
111 Operation.__init__(self, extension)
112 self.command = command
113 self.type = type
115 def __str__(self):
116 return 'Create .%s archive' % self.extension
118 def can_handle(self, data):
119 return isinstance(data, DirData)
121 def save_to_stream(self, data, stream):
122 os.chdir(os.path.dirname(data.path))
123 command = self.command % escape(os.path.basename(data.path))
124 pipe_through_command(command, None, stream)
126 tgz = Extract('tgz', "gunzip -c - | tar xf -")
127 tbz = Extract('tar.bz2', "bunzip2 -c - | tar xf -")
128 rar = Extract('rar', "rar x -")
129 tar = Extract('tar', "tar xf -")
130 rpm = Extract('rpm', "rpm2cpio - | cpio -id --quiet")
131 cpio = Extract('cpio', "cpio -id --quiet")
132 deb = Extract('deb', "ar x '%s'")
133 zip = Extract('zip', "unzip -q '%s'")
134 jar = Extract('jar', "unzip -q '%s'")
136 make_tgz = Archive('tgz', "tar cf - '%s' | gzip", 'application/x-compressed-tar')
137 Archive('tar.gz', "tar cf - '%s' | gzip", 'application/x-compressed-tar')
138 Archive('tar.bz2', "tar cf - '%s' | bzip2", 'application/x-bzip-compressed-tar')
139 Archive('zip', "zip -qr - '%s'", 'application/zip'),
140 Archive('jar', "zip -qr - '%s'", 'application/x-jar')
141 Archive('tar', "tar cf - '%s'", 'application/x-tar')
143 # Note: these go afterwards so that .tar.gz matches before .gz
144 make_gz = Compress('gz', "gzip -c -", 'application/x-gzip')
145 Compress('bz2', "bzip2 -c -", 'application/x-bzip')
147 gz = Decompress('gz', "gunzip -c -")
148 bz2 = Decompress('bz2', "bunzip2 -ck -")
151 # Can bzip2 read bzip files?
153 aliases = {
154 'tar.gz': 'tgz',
155 'tar.bz': 'tar.bz2',
156 'bz': 'bz2'
159 known_extensions = {}
160 for x in operations:
161 try:
162 known_extensions[x.extension] = None
163 except AttributeError:
164 pass
166 class FileData:
167 "A file on the local filesystem."
168 mode = None
169 def __init__(self, path):
170 self.path = path
172 if path == '-':
173 source = sys.stdin
174 else:
175 try:
176 source = file(path)
177 self.mode = os.stat(path).st_mode
178 except:
179 rox.report_exception()
180 sys.exit(1)
182 self.path = path
183 start = source.read(300)
184 try:
185 if source is sys.stdin:
186 raise "Always copy stdin!"
187 source.seek(0)
188 self.source = source
189 except:
190 # Input is not a regular, local, seekable file, so copy it
191 # to a local temp file.
192 import shutil
193 tmp = Tmp()
194 tmp.write(start)
195 tmp.flush()
196 shutil.copyfileobj(source, tmp)
197 self.source = tmp
198 self.source.seek(0)
199 self.default = self.guess_format(start)
201 if path == '-':
202 name = 'Data'
203 else:
204 name = path
205 for ext in known_extensions:
206 if path.endswith('.' + ext):
207 new = path[:-len(ext)-1]
208 if len(new) < len(name):
209 name = new
210 if self.default.add_extension:
211 name += '.' + self.default.extension
212 self.default_name = name
214 def guess_format(self, data):
215 "Return a good default Operation, judging by the first 300 bytes or so."
216 l = len(data)
217 def string(offset, match):
218 return data[offset:offset + len(match)] == match
219 def short(offset, match):
220 if l > offset + 1:
221 a = data[offset]
222 b = data[offset + 1]
223 return ((a == match & 0xff) and (b == (match >> 8))) or \
224 (b == match & 0xff) and (a == (match >> 8))
225 return 0
227 # Archives
228 if string(257, 'ustar\0') or string(257, 'ustar\040\040\0'):
229 return tar
230 if short(0, 070707) or short(0, 0143561) or string(0, '070707') or \
231 string(0, '070701') or string(0, '070702'):
232 return cpio
233 if string(0, '!<arch>') or string(0, '\\<ar>') or string(0, '<ar>'):
234 if string(7, '\ndebian'):
235 return deb
236 if string(0, 'Rar!'): return rar
237 if string(0, 'PK\003\004'): return zip
238 if string(0, '\xed\xab\xee\xdb'): return rpm
240 # Compressed streams
241 if string(0, '\037\213'):
242 if self.path.endswith('.tar.gz') or self.path.endswith('.tgz'):
243 return tgz
244 return gz
245 if string(0, 'BZh') or string(0, 'BZ'):
246 if self.path.endswith('.tar.bz') or self.path.endswith('.tar.bz2') or \
247 self.path.endswith('.tbz') or self.path.endswith('.tbz2'):
248 return tbz
249 return bz2
251 return make_gz
253 class DirData:
254 mode = None
255 def __init__(self, path):
256 self.path = path
257 self.default = make_tgz
258 self.default_name = path + '.' + self.default.extension