Add type detection code.

diff --git a/rarslave-test.py b/rarslave-test.py
index 21aadb8..c57075a 100644 (file)
--- a/rarslave-test.py
+++ b/rarslave-test.py
@@ -46,6 +46,22 @@ class rarslavetest (unittest.TestCase):
 
                self.assertRaises (ValueError, find_all_par2_files, DIR)
 
+       def testHasExtension1 (self):
+               FILE = 'some.file.part01.rar'
+
+               self.assertTrue (has_extension (FILE, 'rar'))
+               self.assertTrue (has_extension (FILE, '.rar'))
+               self.assertTrue (has_extension (FILE, 'part01.rar'))
+               self.assertTrue (has_extension (FILE, '.part01.rar'))
+
+       def testHasExtension2 (self):
+               FILE = 'some.file.part01.rar'
+
+               self.assertFalse (has_extension (FILE, 'zip'))
+               self.assertFalse (has_extension (FILE, '.zip'))
+               self.assertFalse (has_extension (FILE, '.part01'))
+               self.assertFalse (has_extension (FILE, 'part01'))
+
 
 if __name__ == '__main__':
        unittest.main ()

diff --git a/rarslave.py b/rarslave.py
index 7e38816..6532de5 100644 (file)
--- a/rarslave.py
+++ b/rarslave.py
@@ -45,6 +45,73 @@ def find_all_par2_files (dir):
        # Find all files
        return [f for f in os.listdir (dir) if regex.match (f)]
 
+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 (files):
+       """Takes a list of possible files and finds likely heads of
+          extraction."""
+
+       # NOTE: perhaps this should happen AFTER repair is
+       # NOTE: successful. That way all files would already exist
+
+       # According to various sources online:
+       # 1) pre rar-3.0: .rar .r00 .r01 ...
+       # 2) post rar-3.0: .part01.rar .part02.rar 
+       # 3) zip all ver: .zip 
+
+       heads = []
+
+       # Old RAR type, find all files ending in .rar
+       if is_oldrar (files):
+               regex = re.compile ('^.*\.rar$', re.IGNORECASE)
+               for f in files:
+                       if regex.match (f):
+                               heads.append (f)
+
+               return heads
+
+       if is_newrar (files):
+               regex = re.compile ('^.*\.part01.rar$', re.IGNORECASE)
+               for f in files:
+                       if regex.match (f):
+                               heads.append (f)
+
+               return heads
+
+       if is_zip (files):
+               regex = re.compile ('^.*\.zip$', re.IGNORECASE)
+               for f in files:
+                       if regex.match (f):
+                               heads.append (f)
+
+               return heads
+
+       # Not a type we know yet
+       raise ValueError
+
+
+def is_oldrar (files):
+       for f in files:
+               if has_extension (f, '.r00'):
+                       return True
+
+def is_newrar (files):
+       for f in files:
+               if has_extension (f, '.part01.rar'):
+                       return True
+
+def is_zip (files):
+       for f in files:
+               if has_extension (f, '.zip'):
+                       return True
 
 
 def main ():