Fix lint

toltec-dev · Dec 16, 2023 · baa84ae · baa84ae
1 parent f26083b
commit baa84ae
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 84 deletions.
diff --git a/toltec/hooks/patch_rm2fb.py b/toltec/hooks/patch_rm2fb.py
@@ -11,18 +11,14 @@
 import os
 import logging
 import shlex
-import docker
-from elftools.elf.elffile import ELFFile, ELFError
-from toltec import bash
+from elftools.elf.elffile import ELFFile
 from toltec.builder import Builder
 from toltec.recipe import Recipe
 from toltec.util import listener
+from toltec.hooks.strip import walk_elfs, run_in_container, MOUNT_SRC
 
 logger = logging.getLogger(__name__)
 
-MOUNT_SRC = "/src"
-TOOLCHAIN = "toolchain:v1.3.1"
-
 
 def register(builder: Builder) -> None:
     """Register the hook"""
@@ -38,33 +34,17 @@ def post_build(  # pylint: disable=too-many-locals
         # Search for binary objects that can be stripped
         binaries = []
 
-        for directory, _, files in os.walk(src_dir):
-            for file_name in files:
-                file_path = os.path.join(directory, file_name)
-
-                try:
-                    with open(file_path, "rb") as file:
-                        info = ELFFile(file)
-                        symtab = info.get_section_by_name(".symtab")
-
-                        if symtab is None or info.get_machine_arch() != "ARM":
-                            continue
-
-                        dynamic = info.get_section_by_name(".dynamic")
-                        rodata = info.get_section_by_name(".rodata")
-
-                        if (
-                            dynamic
-                            and rodata
-                            and rodata.data().find(b"/dev/fb0") != -1
-                        ):
-                            binaries.append(file_path)
-                except ELFError:
-                    # Ignore non-ELF files
-                    pass
-                except IsADirectoryError:
-                    # Ignore directories
-                    pass
+        def filter_elfs(info: ELFFile) -> None:
+            symtab = info.get_section_by_name(".symtab")
+            if symtab is None or info.get_machine_arch() != "ARM":
+                return
+
+            dynamic = info.get_section_by_name(".dynamic")
+            rodata = info.get_section_by_name(".rodata")
+            if dynamic and rodata and rodata.data().find(b"/dev/fb0") != -1:
+                binaries.append(file_path)
+
+        walk_elfs(src_dir, filter_elfs)
 
         if not binaries:
             logger.debug("Skipping, no arm binaries found")
@@ -94,20 +74,7 @@ def docker_file_path(file_path: str) -> str:
             + " ".join(docker_file_path(file_path) for file_path in binaries)
         )
 
-        logs = bash.run_script_in_container(
-            builder.docker,
-            image=builder.IMAGE_PREFIX + TOOLCHAIN,
-            mounts=[
-                docker.types.Mount(
-                    type="bind",
-                    source=os.path.abspath(src_dir),
-                    target=MOUNT_SRC,
-                )
-            ],
-            variables={},
-            script="\n".join(script),
-        )
-        bash.pipe_logs(logger, logs)
+        run_in_container(builder, src_dir, logger, script)
 
         # Restore original mtimes
         for file_path, mtime in original_mtime.items():

diff --git a/toltec/hooks/strip.py b/toltec/hooks/strip.py
@@ -8,6 +8,7 @@
 import os
 import logging
 import shlex
+from typing import Callable, List
 import docker
 from elftools.elf.elffile import ELFFile, ELFError
 from toltec import bash
@@ -21,6 +22,43 @@
 TOOLCHAIN = "toolchain:v1.3.1"
 
 
+def walk_elfs(src_dir: str, for_each: Callable) -> None:
+    """Walk through all the ELF binaries in a directory and run a method for each of them"""
+    for directory, _, files in os.walk(src_dir):
+        for file_name in files:
+            file_path = os.path.join(directory, file_name)
+
+            try:
+                with open(file_path, "rb") as file:
+                    for_each(ELFFile(file))
+            except ELFError:
+                # Ignore non-ELF files
+                pass
+            except IsADirectoryError:
+                # Ignore directories
+                pass
+
+
+def run_in_container(
+    builder: Builder, src_dir: str, _logger: logging.Logger, script: List[str]
+) -> None:
+    """Run a script in a container and log output"""
+    logs = bash.run_script_in_container(
+        builder.docker,
+        image=builder.IMAGE_PREFIX + TOOLCHAIN,
+        mounts=[
+            docker.types.Mount(
+                type="bind",
+                source=os.path.abspath(src_dir),
+                target=MOUNT_SRC,
+            )
+        ],
+        variables={},
+        script="\n".join(script),
+    )
+    bash.pipe_logs(_logger, logs)
+
+
 def register(builder: Builder) -> None:
     """Register the hook"""
 
@@ -33,29 +71,23 @@ def post_build(  # pylint: disable=too-many-locals,too-many-branches
             return
 
         # Search for binary objects that can be stripped
-        strip_arm = []
-        strip_x86 = []
-
-        for directory, _, files in os.walk(src_dir):
-            for file_name in files:
-                file_path = os.path.join(directory, file_name)
-
-                try:
-                    with open(file_path, "rb") as file:
-                        info = ELFFile(file)
-                        symtab = info.get_section_by_name(".symtab")
-
-                        if symtab:
-                            if info.get_machine_arch() == "ARM":
-                                strip_arm.append(file_path)
-                            elif info.get_machine_arch() in ("x86", "x64"):
-                                strip_x86.append(file_path)
-                except ELFError:
-                    # Ignore non-ELF files
-                    pass
-                except IsADirectoryError:
-                    # Ignore directories
-                    pass
+        strip_arm: List[str] = []
+        strip_x86: List[str] = []
+
+        def filter_elfs(info: ELFFile) -> None:
+            symtab = info.get_section_by_name(".symtab")
+            if not symtab:
+                return
+            if info.get_machine_arch() == "ARM":
+                strip_arm.append(file_path)
+            elif info.get_machine_arch() in ("x86", "x64"):
+                strip_x86.append(file_path)
+
+        walk_elfs(src_dir, filter_elfs)
+
+        if not strip_arm and not strip_x86:
+            logger.debug("Skipping, no binaries found")
+            return
 
         # Save original mtimes to restore them afterwards
         # This will prevent any Makefile rules to be triggered again
@@ -106,20 +138,7 @@ def docker_file_path(file_path: str) -> str:
                     os.path.relpath(file_path, src_dir),
                 )
 
-            logs = bash.run_script_in_container(
-                builder.docker,
-                image=builder.IMAGE_PREFIX + TOOLCHAIN,
-                mounts=[
-                    docker.types.Mount(
-                        type="bind",
-                        source=os.path.abspath(src_dir),
-                        target=MOUNT_SRC,
-                    )
-                ],
-                variables={},
-                script="\n".join(script),
-            )
-            bash.pipe_logs(logger, logs)
+        run_in_container(builder, src_dir, logger, script)
 
         # Restore original mtimes
         for file_path, mtime in original_mtime.items():