extract_utils: move find_file() to utils

Change-Id: I962baf6f8aebfe5073ed75fee2f9e12941bbcfb3

Author: Demon000

extract_utils/extract.py

@@ -6,7 +6,6 @@
 from __future__ import annotations
 
 import os
-import re
 import shutil
 import tarfile
 from os import path
@@ -24,7 +23,13 @@
     ota_extractor_path,
     sdat2img_path,
 )
-from extract_utils.utils import Color, color_print, run_cmd, scan_tree
+from extract_utils.utils import (
+    Color,
+    color_print,
+    find_file,
+    find_files,
+    run_cmd,
+)
 
 ALTERNATE_PARTITION_PATH_MAP = {
     'product': 'system/product',
@@ -102,75 +107,6 @@ def __init__(
         self.extract_all = extract_all
 
 
-def file_name_to_partition(file_name: str):
-    return file_name.split('.', 1)[0]
-
-
-def find_files(
-    input_path: str,
-    partition: Optional[str] = None,
-    name: Optional[str] = None,
-    regex: Optional[str] = None,
-    magic: Optional[bytes] = None,
-    position: int = 0,
-    ext: Optional[str] = None,
-) -> List[str]:
-    file_paths: List[str] = []
-    for file in scan_tree(input_path):
-        if not file.is_file():
-            continue
-
-        file_partition_name = file_name_to_partition(file.name)
-        if partition is not None and partition != file_partition_name:
-            continue
-
-        if name is not None and name != file.name:
-            continue
-
-        if regex is not None and re.match(regex, file.name) is None:
-            continue
-
-        if ext is not None and not file.name.endswith(ext):
-            continue
-
-        if magic is not None:
-            with open(file, 'rb') as f:
-                f.seek(position)
-                file_magic = f.read(len(magic))
-                if file_magic != magic:
-                    continue
-
-        file_paths.append(file.path)
-
-    return file_paths
-
-
-def find_file(
-    input_path: str,
-    partition: Optional[str] = None,
-    name: Optional[str] = None,
-    regex: Optional[str] = None,
-    magic: Optional[bytes] = None,
-    position: int = 0,
-    ext: Optional[str] = None,
-):
-    file_paths = find_files(
-        input_path,
-        partition=partition,
-        name=name,
-        regex=regex,
-        magic=magic,
-        position=position,
-        ext=ext,
-    )
-
-    assert len(file_paths) <= 1
-    if file_paths:
-        return file_paths[0]
-
-    return None
-
-
 def find_alternate_partitions(
     extract_partitions: List[str],
     found_partitions: Iterable[str],

extract_utils/utils.py

@@ -8,6 +8,7 @@
 import hashlib
 import importlib.util
 import os
+import re
 import shutil
 from contextlib import contextmanager
 from enum import Enum
@@ -305,3 +306,72 @@ def read_mmap_chunked(
 def write_zero(f: BinaryIO, size: int):
     f.seek(size - 1, SEEK_CUR)
     f.write(b'\x00')
+
+
+def file_name_to_partition(file_name: str):
+    return file_name.split('.', 1)[0]
+
+
+def find_files(
+    input_path: str,
+    partition: Optional[str] = None,
+    name: Optional[str] = None,
+    regex: Optional[str] = None,
+    magic: Optional[bytes] = None,
+    position: int = 0,
+    ext: Optional[str] = None,
+) -> List[str]:
+    file_paths: List[str] = []
+    for file in scan_tree(input_path):
+        if not file.is_file():
+            continue
+
+        file_partition_name = file_name_to_partition(file.name)
+        if partition is not None and partition != file_partition_name:
+            continue
+
+        if name is not None and name != file.name:
+            continue
+
+        if regex is not None and re.match(regex, file.name) is None:
+            continue
+
+        if ext is not None and not file.name.endswith(ext):
+            continue
+
+        if magic is not None:
+            with open(file, 'rb') as f:
+                f.seek(position)
+                file_magic = f.read(len(magic))
+                if file_magic != magic:
+                    continue
+
+        file_paths.append(file.path)
+
+    return file_paths
+
+
+def find_file(
+    input_path: str,
+    partition: Optional[str] = None,
+    name: Optional[str] = None,
+    regex: Optional[str] = None,
+    magic: Optional[bytes] = None,
+    position: int = 0,
+    ext: Optional[str] = None,
+):
+    file_paths = find_files(
+        input_path,
+        partition=partition,
+        name=name,
+        regex=regex,
+        magic=magic,
+        position=position,
+        ext=ext,
+    )
+
+    assert len(file_paths) <= 1
+    if file_paths:
+        return file_paths[0]
+
+    return None