Change the directory walker to a more sane implementation
[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 import logging
39 from optparse import OptionParser
40
41 ### Default Configuration Variables ###
42 DICT_FILE = os.path.join ('~','.config','animesorter2','animesorter.dict')
43 WORK_DIR =  os.path.join ('~','downloads','usenet')
44 SORT_DIR =  os.path.join ('/','data','Anime')
45 TYPES_REGEX = '.*(avi|ogm|mkv|mp4|\d\d\d)$'
46
47
48 class AnimeSorter2:
49
50     def __init__(self, options):
51         self.options = options
52         self.dict = None
53
54     def __valid_dict_line (self, line):
55         if len(line) <= 0:
56             return False
57
58         if '=' not in line:
59             return False
60
61         # Comment lines are not really valid
62         if re.match ('^(\s*#.*|\s*)$', line):
63             return False
64
65         # Make sure that there is a definition and it is valid
66         try:
67             (regex, directory) = line.split('=')
68             regex = regex.strip()
69             directory = directory.strip()
70         except:
71             return False
72
73         # Make sure they have length
74         if len(regex) <= 0 or len(directory) <= 0:
75             return False
76
77         # I guess that it's valid now
78         return True
79
80     def parse_dict(self):
81         """Parses a dictionary file containing the sort definitions in the form:
82         REGEX_PATTERN = DIRECTORY
83
84         Returns a list of tuples of the form (compiled_regex, to_directory)"""
85
86         try:
87             f = open(self.options.dict_file, 'r', 0)
88             try:
89                 raw_lines = f.readlines()
90             finally:
91                 f.close()
92         except IOError:
93             logging.critical ('Opening dictionary: %s FAILED' % self.options.dict_file)
94             sys.exit()
95
96         ### Find all of the valid lines in the file
97         valid_lines = [l for l in raw_lines if self.__valid_dict_line (l)]
98
99         # Set up variable for result
100         result = []
101
102         ### Split each line into a tuple, and strip each element of spaces
103         for l in valid_lines:
104             (regex, directory) = l.split('=')
105             regex = regex.strip()
106             directory = directory.strip()
107
108             # Fix up the directory if necessary
109             if directory[0] != '/':
110                 directory = os.path.join (self.options.output_dir, directory)
111
112             # Fix up the regex
113             if regex[0] != '^':
114                 regex = '^' + regex
115
116             if regex[-1] != '$':
117                 regex += '$'
118
119             # Store the result
120             result.append ( (re.compile (regex), directory) )
121
122         ### Give some information about the dictionary we are using
123         logging.info ('Successfully loaded %d records from %s\n' % \
124                 (len(result), self.options.dict_file))
125
126         return tuple (result)
127
128     def as_makedirs (self, dirname):
129         """Call os.makedirs(dirname), but check first whether we are in pretend
130            mode, or if we're running interactively."""
131
132         if not os.path.isdir (dirname):
133
134             if self.options.pretend:
135                 logging.info ('Will create directory %s' % dirname)
136                 return 0
137
138             if self.get_user_choice ('Make directory?: %s' % (dirname, )):
139
140                 try:
141                     os.makedirs (dirname)
142                     logging.info ('Created directory %s' % dirname)
143                 except:
144                     logging.critical ('Failed to create directory %s' % dirname)
145                     return errno.EIO
146
147         return 0
148
149     def as_move_single_file (self, f, fromdir, todir):
150         """Move the single file named $f from the directory $fromdir to the
151            directory $todir"""
152
153         srcname = os.path.join (fromdir, f)
154         dstname = os.path.join (todir, f)
155
156         if self.options.pretend:
157             logging.info ('Will move %s to %s' % (f, todir))
158             return 0 # success
159
160         if self.get_user_choice ('Move file?: %s --> %s' % (srcname, dstname)):
161             try:
162                 shutil.move (srcname, dstname)
163                 logging.info ('Moved %s to %s' % (f, todir))
164             except:
165                 logging.critical ('FAILED to move %s to %s' % (f, todir))
166                 return errno.EIO
167
168         return 0
169
170     def move_files(self, files, fromdir, todir):
171         """move_files(files, fromdir, todir):
172         Move the files represented by the list FILES from FROMDIR to TODIR"""
173
174         ret = 0
175
176         # Leave immediately if we have nothing to do
177         if len(files) <= 0:
178             return ret
179
180         ## Create the directory if it doesn't exist
181         ret = self.as_makedirs (todir)
182
183         if ret:
184             # we cannot continue, since we can't make the directory
185             return ret
186
187         ## Try to move every file, one at a time
188         for f in files:
189             ret = self.as_move_single_file (f, fromdir, todir)
190
191             if ret:
192                 # something bad happened when moving a file
193                 break
194
195         return ret
196
197     def __dir_walker(self, rootdir, files):
198
199         for (r,d) in self.dict:
200             matches = [f for f in files if r.match(f)]
201             self.move_files (matches, rootdir, d)
202
203     def get_user_choice(self, prompt):
204
205         # If we're not in interactive mode, then always return True
206         if self.options.interactive == False:
207             return True
208
209         # Get the user's choice since we're not in interactive mode
210         done = False
211         while not done:
212             s = raw_input('%s [y/N]: ' % (prompt, )).lower()
213
214             if s == 'y' or s == 'yes':
215                 return True
216
217             if s == 'n' or s == 'no' or s == '':
218                 return False
219
220             print 'Response not understood, try again.'
221
222     def main(self):
223
224         ## Print the program's header
225         logging.info ('Regular Expression File Sorter (aka animesorter)')
226         logging.info ('=' * 80)
227         logging.info ('Copyright (c) 2005-2007, Ira W. Snyder (devel@irasnyder.com)')
228         logging.info ('This program is licensed under the GNU GPL v2')
229         logging.info ('')
230
231         ## Parse the dictionary
232         self.dict = self.parse_dict()
233
234         if self.options.recursive:
235             ## Start walking through directories
236             for root, dirs, files in os.walk(self.options.start_dir):
237                 self.__dir_walker(root, files)
238         else:
239             self.__dir_walker(self.options.start_dir, [f for f in
240                     os.listdir(self.options.start_dir) if os.path.isfile(f)])
241
242
243 ### MAIN IS HERE ###
244 def main():
245
246     # Set up the logger
247     logging.basicConfig (level=logging.INFO, format='%(message)s')
248
249     ### Get the program options
250     parser = OptionParser()
251     parser.add_option('-q', '--quiet', action='store_true', dest='quiet',
252             default=False, help="Don't print status messages to stdout")
253     parser.add_option('-d', '--dict', dest='dict_file', default=DICT_FILE,
254             help='Read dictionary from FILE', metavar='FILE')
255     parser.add_option('-n', '--not-recursive', action='store_false', dest='recursive',
256             default=True, help='don\'t run recursively')
257     parser.add_option('-s', '--start-dir', dest='start_dir', default=WORK_DIR,
258             help='Start running at directory DIR', metavar='DIR')
259     parser.add_option('-o', '--output-dir', dest='output_dir', default=SORT_DIR,
260             help='Sort files into DIR', metavar='DIR')
261     parser.add_option('-i', '--interactive', dest='interactive', default=False,
262             help='Confirm each move', action='store_true')
263     parser.add_option('-p', '--pretend', dest='pretend', default=False,
264             help='Enable pretend mode', action='store_true')
265
266     ## Parse the options
267     (options, args) = parser.parse_args()
268
269     ## Correct directories
270     options.dict_file = os.path.abspath(os.path.expanduser(options.dict_file))
271     options.start_dir = os.path.abspath(os.path.expanduser(options.start_dir))
272     options.output_dir = os.path.abspath(os.path.expanduser(options.output_dir))
273
274     # Change the loglevel if we're running in quiet mode
275     if options.quiet:
276         logging.getLogger().setLevel (logging.CRITICAL)
277
278     as = AnimeSorter2(options)
279     as.main()
280
281 if __name__ == '__main__':
282     main ()
283