refactor: get_extract_dir_for_input -> ExtractionConfig.get_extract_dir_for

Author: e3krisztian

tests/test_processing.py

@@ -10,7 +10,6 @@
     calculate_entropy,
     calculate_unknown_chunks,
     draw_entropy_plot,
-    get_extract_dir_for_input,
     remove_inner_chunks,
 )
 
@@ -155,10 +154,10 @@ def test_calculate_entropy_no_exception(path: Path, draw_plot: bool):
         ("/extract", "/some/place/else/firmware", "firmware"),
     ],
 )
-def test_get_extract_dir_for_input(
+def test_ExtractionConfig_get_extract_dir_for(
     extract_root: str, path: str, extract_dir_prefix: str
 ):
     cfg = ExtractionConfig(extract_root=Path(extract_root), entropy_depth=0)
-    assert get_extract_dir_for_input(cfg, Path(path)) == (
+    assert cfg.get_extract_dir_for(Path(path)) == (
         cfg.extract_root / Path(extract_dir_prefix + cfg.extract_suffix)
     )

unblob/processing.py

@@ -65,6 +65,17 @@ class ExtractionConfig:
     extract_suffix: str = "_extract"
     handlers: Handlers = BUILTIN_HANDLERS
 
+    def get_extract_dir_for(self, path: Path) -> Path:
+        """Extraction dir under root with the name of path."""
+        try:
+            relative_path = path.relative_to(self.extract_root)
+        except ValueError:
+            # path is not inside root, i.e. it is an input file
+            relative_path = Path(path.name)
+        extract_name = path.name + self.extract_suffix
+        extract_dir = self.extract_root / relative_path.with_name(extract_name)
+        return extract_dir.expanduser().resolve()
+
 
 @terminate_gracefully
 def process_file(config: ExtractionConfig, path: Path) -> ProcessResult:
@@ -104,7 +115,7 @@ def process_result(pool, result):
 def check_extract_directory(task: Task, config: ExtractionConfig):
     errors = []
 
-    extract_dir = get_extract_dir_for_input(config, task.path)
+    extract_dir = config.get_extract_dir_for(task.path)
     if extract_dir.exists():
         if config.force_extract:
             shutil.rmtree(extract_dir)
@@ -198,7 +209,7 @@ def __init__(
         self.size = size
         self.result = result
 
-        self.carve_dir = get_extract_dir_for_input(config, self.task.path)
+        self.carve_dir = config.get_extract_dir_for(self.task.path)
 
     def process(self):
         logger.debug("Processing file", path=self.task.path, size=self.size)
@@ -285,18 +296,6 @@ def _extract_chunk(self, file, chunk: ValidChunk):
             )
 
 
-def get_extract_dir_for_input(config: ExtractionConfig, path: Path) -> Path:
-    """Extraction dir under root with the name of path."""
-    try:
-        relative_path = path.relative_to(config.extract_root)
-    except ValueError:
-        # path is not inside root, i.e. it is an input file
-        relative_path = Path(path.name)
-    extract_name = path.name + config.extract_suffix
-    extract_dir = config.extract_root / relative_path.with_name(extract_name)
-    return extract_dir.expanduser().resolve()
-
-
 def remove_inner_chunks(chunks: List[ValidChunk]) -> List[ValidChunk]:
     """Remove all chunks from the list which are within another bigger chunks."""
     if not chunks: