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