[RARSLAVE] Improve "likely file" detection

[rarslave2.git] / rarslave.py

diff --git a/rarslave.py b/rarslave.py
index 5be2b2f..1969ce2 100644 (file)
--- a/rarslave.py
+++ b/rarslave.py
@@ -134,7 +134,7 @@ class RarslaveRepairer (object):
 
                # Get set up
                basename = get_basename (self.file)
-               all_files = find_likely_files (basename, self.dir)
+               all_files = find_likely_files (self.dir, self.file)
                all_files.sort ()
                par2_files = find_par2_files (all_files)
 
@@ -194,18 +194,23 @@ def get_basename (name):
 
        return name
 
-def find_likely_files (name, dir):
+def find_likely_files (dir, p2file):
        """Finds files which are likely to be part of the set corresponding
           to $name in the directory $dir"""
 
-       if not os.path.isdir (os.path.abspath (dir)):
-               raise ValueError # bad directory given
+       assert os.path.isdir (dir)
+       assert os.path.isfile (os.path.join (dir, p2file))
+
+       basename = get_basename (p2file)
 
        dir = os.path.abspath (dir)
-       ename = re.escape (name)
+       ename = re.escape (basename)
        regex = re.compile ('^%s.*$' % (ename, ))
 
-       return [f for f in os.listdir (dir) if regex.match (f)]
+       name_matches = [f for f in os.listdir (dir) if regex.match (f)]
+       parsed_matches = par2parser.get_protected_files (dir, p2file)
+
+       return name_matches + parsed_matches
 
 def find_par2_files (files):
        """Find all par2 files in the list $files"""
@@ -226,16 +231,6 @@ def find_all_par2_files (dir):
 
        return find_par2_files (files)
 
-def has_extension (f, ext):
-       """Checks if f has the extension ext"""
-
-       if ext[0] != '.':
-               ext = '.' + ext
-
-       ext = re.escape (ext)
-       regex = re.compile ('^.*%s$' % (ext, ), re.IGNORECASE)
-       return regex.match (f)
-
 def find_extraction_heads (dir, files):
        """Takes a list of possible files and finds likely heads of
           extraction."""
@@ -261,7 +256,7 @@ def find_extraction_heads (dir, files):
 
        if is_newrar (files):
                extractor = RarslaveExtractor (TYPE_NEWRAR)
-               regex = re.compile ('^.*\.part01.rar$', re.IGNORECASE)
+               regex = re.compile ('^.*\.part0*1.rar$', re.IGNORECASE)
                for f in files:
                        if regex.match (f):
                                extractor.addHead (dir, f)
@@ -307,35 +302,34 @@ def find_extraction_heads (dir, files):
 
        return extractor
 
-def is_oldrar (files):
-       for f in files:
-               if has_extension (f, '.r00'):
-                       return True
+def generic_matcher (files, regex, nocase=False):
+       """Run the regex over the files, and see if one matches or not.
+       NOTE: this does not return the matches, just if a match occurred."""
 
-       return False
+       if nocase:
+               cregex = re.compile (regex, re.IGNORECASE)
+       else:
+               cregex = re.compile (regex)
 
-def is_newrar (files):
        for f in files:
-               if has_extension (f, '.part01.rar'):
+               if cregex.match (f):
                        return True
 
        return False
 
-def is_zip (files):
-       for f in files:
-               if has_extension (f, '.zip'):
-                       return True
+def is_oldrar (files):
+       return generic_matcher (files, '^.*\.r00$')
 
-       return False
+def is_newrar (files):
+       return generic_matcher (files, '^.*\.part0*1\.rar$')
+
+def is_zip (files):
+       return generic_matcher (files, '^.*\.zip$')
 
 def is_noextract (files):
        # Type that needs no extraction.
        # TODO: Add others ???
-       for f in files:
-               if has_extension (f, '.001'):
-                       return True
-
-       return False
+       return generic_matcher (files, '^.*\.001$')
 
 def find_deleteable_files (files):
        # Deleteable types regex should come from the config
@@ -363,7 +357,7 @@ class PAR2Set (object):
                self.file = file
 
                basename = get_basename (file)
-               self.likely_files = find_likely_files (basename, dir)
+               self.likely_files = find_likely_files (dir, file)
 
        def __list_eq (self, l1, l2):
 
@@ -405,7 +399,7 @@ class PAR2Set (object):
                # Deletion Stage
                DELETE_INTERACTIVE = options.interactive
                deleteable_files = find_deleteable_files (self.likely_files)
-               ret = delete_list (deleteable_files, DELETE_INTERACTIVE)
+               ret = delete_list (self.dir, deleteable_files, DELETE_INTERACTIVE)
 
                if ret != SUCCESS:
                        logger.addMessage ('Deletion stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
@@ -414,9 +408,11 @@ class PAR2Set (object):
                logger.addMessage ('Successfully completed: %s' % par2head)
                return SUCCESS
 
-def delete_list (files, interactive=False):
+def delete_list (dir, files, interactive=False):
        # Delete a list of files
 
+       assert os.path.isdir (dir)
+
        done = False
        valid_y = ['Y', 'YES']
        valid_n = ['N', 'NO']
@@ -433,7 +429,7 @@ def delete_list (files, interactive=False):
                        return SUCCESS
 
        for f in files:
-               os.remove (f)
+               os.remove (os.path.join (dir, f))
 
        return SUCCESS