Add copyright notice.
[animesorter.git] / animesorter2.py
1 #!/usr/bin/env python
2 # vim: set ts=4 sts=4 sw=4 textwidth=92 expandtab:
3
4 """
5 The animesorter program
6
7 This will sort arbitrary files into directories by regular expression.
8 """
9
10 __author__    = "Ira W. Snyder (devel@irasnyder.com)"
11 __copyright__ = "Copyright (c) 2006,2007 Ira W. Snyder (devel@irasnyder.com)"
12 __license__   = "GNU GPL v2 (or, at your option, any later version)"
13
14 #    animesorter2.py -- a regular expression file-sorting utility
15 #
16 #    Copyright (C) 2006,2007  Ira W. Snyder (devel@irasnyder.com)
17 #
18 #    This program is free software; you can redistribute it and/or modify
19 #    it under the terms of the GNU General Public License as published by
20 #    the Free Software Foundation; either version 2 of the License, or
21 #    (at your option) any later version.
22 #
23 #    This program is distributed in the hope that it will be useful,
24 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
25 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
26 #    GNU General Public License for more details.
27 #
28 #    You should have received a copy of the GNU General Public License
29 #    along with this program; if not, write to the Free Software
30 #    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
31
32
33 import os
34 import re
35 import sys
36 import errno
37 import shutil
38 from optparse import OptionParser
39
40 ### Default Configuration Variables ###
41 DICT_FILE = os.path.join ('~','.config','animesorter2','animesorter.dict')
42 WORK_DIR =  os.path.join ('~','downloads','usenet')
43 SORT_DIR =  os.path.join ('/','data','Anime')
44 TYPES_REGEX = '.*(avi|ogm|mkv|mp4|\d\d\d)$'
45
46
47 class AnimeSorter2:
48
49     def __init__(self, options):
50         self.options = options
51
52     def parse_dict(self):
53         """Parses a dictionary file containing the sort definitions in the form:
54         DIRECTORY = PATTERN
55
56         Returns a list of tuples of the form (compiled_regex, to_directory)"""
57
58         try:
59             f = open(self.options.dict_file, 'r', 0)
60             try:
61                 data = f.read()
62             finally:
63                 f.close()
64         except IOError:
65             self.print_dict_fail (self.options.dict_file)
66             sys.exit()
67
68         ### Get a LIST containing each line in the file
69         lines = [l for l in data.split('\n') if len(l) > 0]
70
71         ### Remove comments / blank lines (zero length lines already removed above)
72         regex = re.compile ('^\s*#.*$')
73         lines = [l for l in lines if not re.match (regex, l)]
74         regex = re.compile ('^\s*$')
75         lines = [l for l in lines if not re.match (regex, l)]
76
77         ### Split each line into a tuple, and strip each element of spaces
78         result = self.split_lines(lines)
79         result = [(re.compile(r), d) for r, d in result]
80
81         ### Give some information about the dictionary we are using
82         self.print_dict_suc (self.options.dict_file, len(result))
83
84         return tuple(result)
85
86     def split_lines(self, lines):
87
88         result = []
89
90         for l in lines:
91
92             try:
93                 r, d = l.split('=')
94                 r = r.strip()
95                 d = d.strip()
96             except ValueError:
97                 self.print_dict_bad_line (l)
98                 continue
99
100             result.append((r, d))
101
102         return result
103
104     def get_matches(self, files, pattern):
105         """get_matches(files, pattern):
106
107         files is type LIST
108         pattern is type sre.SRE_Pattern
109
110         Returns a list of the files matching the pattern as type sre.SRE_Match."""
111
112         matches = [m for m in files if pattern.search(m)]
113         return matches
114
115     def as_makedirs (self, dirname):
116         """Call os.makedirs(dirname), but check first whether we are in pretend
117            mode, or if we're running interactively."""
118
119         if not os.path.isdir (dirname):
120
121             if self.options.pretend:
122                 self.print_dir_create_pretend (dirname)
123                 return 0
124
125             if self.get_user_choice ('Make directory?: %s' % (dirname, )):
126
127                 try:
128                     os.makedirs (dirname)
129                     self.print_dir_create_suc (dirname)
130                 except:
131                     self.print_dir_create_fail (dirname)
132                     return errno.EIO
133
134         return 0
135
136     def as_move_single_file (self, f, fromdir, todir):
137         """Move the single file named $f from the directory $fromdir to the
138            directory $todir"""
139
140         srcname = os.path.join (fromdir, f)
141         dstname = os.path.join (todir, f)
142
143         if self.options.pretend:
144             self.print_move_file_pretend (f, todir)
145             return 0 # success
146
147         if self.get_user_choice ('Move file?: %s --> %s' % (srcname, dstname)):
148             try:
149                 shutil.move (srcname, dstname)
150                 self.print_move_file_suc (f, todir)
151             except:
152                 self.print_move_file_fail (f, todir)
153                 return errno.EIO
154
155         return 0
156
157     def move_files(self, files, fromdir, todir):
158         """move_files(files, fromdir, todir):
159         Move the files represented by the list FILES from FROMDIR to TODIR"""
160
161         ret = 0
162
163         ## Check for a non-default directory
164         if todir[0] != '/':
165             todir = os.path.join(self.options.output_dir, todir)
166
167         ## Create the directory if it doesn't exist
168         ret = self.as_makedirs (todir)
169
170         if ret:
171             # we cannot continue, since we can't make the directory
172             return ret
173
174         ## Try to move every file, one at a time
175         for f in files:
176             ret = self.as_move_single_file (f, fromdir, todir)
177
178             if ret:
179                 # something bad happened when moving a file
180                 break
181
182         return ret
183
184     def __dir_walker(self, dict, root, dirs, files):
185
186         ## Get all of the files in the directory that are of the correct types
187         types_re = re.compile(TYPES_REGEX, re.IGNORECASE)
188         raw_matches = [f for f in files if types_re.match(f)]
189
190         ### Loop through the dictionary and try to move everything that matches
191         for regex, todir in dict:
192             matches = self.get_matches(raw_matches, regex)
193
194             ## Move the files if we've found some
195             if len(matches) > 0:
196                 self.move_files(matches, root, todir)
197
198     def get_user_choice(self, prompt):
199
200         # If we're not in interactive mode, then always return True
201         if self.options.interactive == False:
202             return True
203
204         # Get the user's choice since we're not in interactive mode
205         done = False
206         while not done:
207             s = raw_input('%s [y/N]: ' % (prompt, )).lower()
208
209             if s == 'y' or s == 'yes':
210                 return True
211
212             if s == 'n' or s == 'no' or s == '':
213                 return False
214
215             print 'Response not understood, try again.'
216
217     def main(self):
218
219         ## Print the program's header
220         self.print_prog_header ()
221
222         ## Parse the dictionary
223         dict = self.parse_dict()
224
225         if self.options.recursive:
226             ## Start walking through directories
227             for root, dirs, files in os.walk(self.options.start_dir):
228                 self.__dir_walker(dict, root, dirs, files)
229         else:
230             self.__dir_walker(dict, self.options.start_dir,
231                     [d for d in os.listdir(self.options.start_dir) if os.path.isdir(d)],
232                     [f for f in os.listdir(self.options.start_dir) if os.path.isfile(f)])
233
234     ############################################################################
235     ### Functions for the printing system
236     ############################################################################
237
238     def print_prog_header(self):
239         if self.options.quiet:
240             return
241
242         print 'Regular Expression File Sorter (aka animesorter)'
243         print '================================================================================'
244         print 'Copyright (c) 2005,2006, Ira W. Snyder (devel@irasnyder.com)'
245         print 'All rights reserved.'
246         print 'This program is licensed under the GNU GPL v2'
247         print
248
249     def print_move_file_suc(self, f, t):
250         if self.options.quiet:
251             return
252
253         print 'Moved %s to %s' % (f, t)
254
255     def print_move_file_fail(self, f, t):
256         print 'FAILED to move %s to %s' % (f, t)
257
258     def print_move_file_pretend (self, f, t):
259         print 'Will move %s to %s' % (f, t)
260
261     def print_dir_create_suc(self, d):
262         if self.options.quiet:
263             return
264
265         print 'Created directory %s' % (d, )
266
267     def print_dir_create_fail(self, d):
268         print 'Failed to create directory %s' % (d, )
269
270     def print_dir_create_pretend (self, d):
271         print 'Will create directory %s' % (d, )
272
273     def print_dict_suc(self, dic, num):
274         if self.options.quiet:
275             return
276
277         print 'Successfully loaded %d records from %s\n' % (num, dic)
278
279     def print_dict_fail(self, dic):
280         print 'Opening dictionary: %s FAILED' % (dic, )
281
282     def print_dict_bad_line(self, dic):
283         print 'Bad line in dictionary: %s' % (dic, )
284
285
286
287 ### MAIN IS HERE ###
288 def main():
289
290     ### Get the program options
291     parser = OptionParser()
292     parser.add_option('-q', '--quiet', action='store_true', dest='quiet',
293             default=False, help="Don't print status messages to stdout")
294     parser.add_option('-d', '--dict', dest='dict_file', default=DICT_FILE,
295             help='Read dictionary from FILE', metavar='FILE')
296     parser.add_option('-n', '--not-recursive', action='store_false', dest='recursive',
297             default=True, help='don\'t run recursively')
298     parser.add_option('-s', '--start-dir', dest='start_dir', default=WORK_DIR,
299             help='Start running at directory DIR', metavar='DIR')
300     parser.add_option('-o', '--output-dir', dest='output_dir', default=SORT_DIR,
301             help='Sort files into DIR', metavar='DIR')
302     parser.add_option('-i', '--interactive', dest='interactive', default=False,
303             help='Confirm each move', action='store_true')
304     parser.add_option('-p', '--pretend', dest='pretend', default=False,
305             help='Enable pretend mode', action='store_true')
306
307     ## Parse the options
308     (options, args) = parser.parse_args()
309
310     ## Correct directories
311     options.dict_file = os.path.abspath(os.path.expanduser(options.dict_file))
312     options.start_dir = os.path.abspath(os.path.expanduser(options.start_dir))
313     options.output_dir = os.path.abspath(os.path.expanduser(options.output_dir))
314
315     as = AnimeSorter2(options)
316     as.main()
317
318 if __name__ == '__main__':
319     main ()
320
321 # vim: set ts=4 sw=4 sts=4 expandtab: