Optimize find_backup_data to take O(n) instead of O(n^2)

tibi_hardlinks/backups.py

@@ -12,6 +12,11 @@ HEIRACHY = (
     "app_version_name",
     "backup_time",
 )
+PROP_FILE_REGEX = re.compile(
+    r"(?P<apk_name>(?:\w+\.)*\w+)-(?P<date>\d{8})-(?P<time>\d{6})"
+)
+
+COMPRESSION_SUFFIX = {"NONE": "", "BZIP2": ".bz2", "GZIP": ".gz"}
 
 
 def create_backup_directory_name(backup_info, time_format="%Y-%m-%d_%H-%M-%S"):
@@ -39,28 +44,49 @@ def create_backup_directory_name(backup_info, time_format="%Y-%m-%d_%H-%M-%S"):
     return backup_directory_name
 
 
+def _find_backup_data(property_file: Path, backup_info, compression_suffix):
+    backup_root = property_file.parent
+    parsed_property_file = PROP_FILE_REGEX.match(property_file.name)
+
+    output = {"property_file": property_file, "apk": None, "data": None}
+    if parsed_property_file:
+        apk_name = parsed_property_file.group("apk_name")
+        # apk
+        try:
+            apk_hash = backup_info["app_apk_md5"]
+            apk_path = backup_root / f"{apk_name}-{apk_hash}.apk{compression_suffix}"
+            if apk_path.exists():
+                output["apk"] = apk_path
+        except KeyError:
+            pass
+        # tar
+        tar_path = backup_root / f"{property_file.stem}.tar{compression_suffix}"
+
+        if tar_path.exists():
+            output["data"] = tar_path
+        else:
+            # XML
+            xml_path = backup_root / f"{property_file.stem}.xml{compression_suffix}"
+            if xml_path.exists():
+                output["data"] = xml_path
+    else:
+        raise NotImplementedError(f"Unexpected behavior for {property_file}")
+    return output
+
+
 def find_backup_data(property_file: Path, backup_info):
     # find the path of the *.apk, *.properties, and *.tar(.gz) corresponding to property_file
-    backup_root = property_file.parent
+    compression = None
     try:
-        apk_hash = backup_info["app_apk_md5"]
-        apk_file = next(
-            filter(
-                lambda apk_file: apk_hash in apk_file.name, backup_root.glob("*.apk*")
-            )
-        )
-    except (KeyError, StopIteration):
-        apk_file = None
-    related_files = list(backup_root.glob(property_file.stem + "*"))
-    if len(related_files) > 2:
+        compression = backup_info["app_apk_codec"]
+    except KeyError:
+        for compression_suffix in COMPRESSION_SUFFIX.values():
+            output = _find_backup_data(property_file, backup_info, compression_suffix)
+            if not (output["apk"] is None and output["data"] is None):
+                return output
         raise Exception
-    tar_file = None
-    for filename in related_files:
-        if ".tar" in filename.name:
-            tar_file = filename
-
-    property_file = property_file
-    return {"property_file": property_file, "apk_file": apk_file, "tar_file": tar_file}
+    compression_suffix = COMPRESSION_SUFFIX[compression]
+    return _find_backup_data(property_file, backup_info, compression_suffix)
 
 
 class Backup: