jumpstarter_driver_flashers: fmt

Author: NickCao

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/bundle.py

@@ -9,27 +9,26 @@ class FileAddress(BaseModel):
     file: str
     address: str
 
+
 class DtbVariant(BaseModel):
     default: str
     address: str
     variants: dict[str, str]
 
+
 class FlasherLogin(BaseModel):
     login_prompt: str
     username: str | None = None
     password: str | None = None
     prompt: str
 
+
 class FlashBundleSpecV1Alpha1(BaseModel):
     manufacturer: str
     link: Optional[str]
     bootcmd: str
     shelltype: Literal["busybox"] = Field(default="busybox")
-    login: FlasherLogin = Field(
-        default_factory=lambda: FlasherLogin(
-            login_prompt="login:",
-            prompt="#")
-        )
+    login: FlasherLogin = Field(default_factory=lambda: FlasherLogin(login_prompt="login:", prompt="#"))
     default_target: str
     targets: dict[str, str]
     kernel: FileAddress
@@ -41,6 +40,7 @@ class FlashBundleSpecV1Alpha1(BaseModel):
 class ObjectMeta(BaseModel):
     name: str
 
+
 class FlasherBundleManifestV1Alpha1(BaseModel):
     apiVersion: Literal["jumpstarter.dev/v1alpha1"] = Field(default="jumpstarter.dev/v1alpha1")
     kind: Literal["FlashBundleManifest"] = Field(default="FlashBundleManifest")

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py

@@ -25,6 +25,7 @@
 
 debug_console_option = click.option("--console-debug", is_flag=True, help="Enable console debug mode")
 
+
 @dataclass(kw_only=True)
 class BaseFlasherClient(FlasherClient, CompositeClient):
     """
@@ -99,10 +100,11 @@ def flash(
             error_queue = Queue()
 
             # Start the storage write operation in the background
-            storage_thread = threading.Thread(target=self._transfer_bg_thread,
-                                            args=(path, operator, operator_scheme,
-                                                  os_image_checksum, self.http.storage, error_queue),
-                                            name="storage_transfer")
+            storage_thread = threading.Thread(
+                target=self._transfer_bg_thread,
+                args=(path, operator, operator_scheme, os_image_checksum, self.http.storage, error_queue),
+                name="storage_transfer",
+            )
             storage_thread.start()
 
         # Make the exporter download the bundle contents and set files in the right places
@@ -155,7 +157,6 @@ def flash(
                 self.logger.info("Powering off target")
                 self.power.off()
 
-
     def _flash_with_progress(self, console, manifest, path, image_url, target_path):
         """Flash image to target device with progress monitoring.
 
@@ -170,8 +171,8 @@ def _flash_with_progress(self, console, manifest, path, image_url, target_path):
         decompress_cmd = _get_decompression_command(path)
         flash_cmd = (
             f'( wget -q -O - "{image_url}" | '
-            f'{decompress_cmd} '
-            f'dd of={target_path} bs=64k iflag=fullblock oflag=direct) &'
+            f"{decompress_cmd} "
+            f"dd of={target_path} bs=64k iflag=fullblock oflag=direct) &"
         )
         console.sendline(flash_cmd)
         console.expect(manifest.spec.login.prompt, timeout=60)
@@ -190,16 +191,16 @@ def _flash_with_progress(self, console, manifest, path, image_url, target_path):
             if "No such file or directory" in console.before.decode(errors="ignore"):
                 break
             data = console.before.decode(errors="ignore")
-            match = re.search(r'pos:\s+(\d+)', data)
+            match = re.search(r"pos:\s+(\d+)", data)
             if match:
                 current_bytes = int(match.group(1))
                 current_time = time.time()
                 elapsed = current_time - last_time
 
                 if elapsed >= 1.0:  # Update speed every second
                     bytes_diff = current_bytes - last_pos
-                    speed_mb = (bytes_diff / (1024*1024)) / elapsed
-                    total_mb = current_bytes/(1024*1024)
+                    speed_mb = (bytes_diff / (1024 * 1024)) / elapsed
+                    total_mb = current_bytes / (1024 * 1024)
                     self.logger.info(f"Flash progress: {total_mb:.2f} MB, Speed: {speed_mb:.2f} MB/s")
 
                     last_pos = current_bytes
@@ -209,7 +210,6 @@ def _flash_with_progress(self, console, manifest, path, image_url, target_path):
         console.sendline("sync")
         console.expect(manifest.spec.login.prompt, timeout=1200)
 
-
     def _get_target_device(self, target: str, manifest: FlasherBundleManifestV1Alpha1, console) -> str:
         """Get the target device path from the manifest, resolving block devices if needed.
 
@@ -229,15 +229,19 @@ def _get_target_device(self, target: str, manifest: FlasherBundleManifestV1Alpha
             raise ArgumentError(f"Target {target} not found in manifest")
 
         if target_path.startswith("/sys/class/block#"):
-            target_path = self._lookup_block_device(
-                console, manifest.spec.login.prompt, target_path.split("#")[1])
+            target_path = self._lookup_block_device(console, manifest.spec.login.prompt, target_path.split("#")[1])
 
         return target_path
 
-
-    def _transfer_bg_thread(self, src_path: PathBuf, src_operator: Operator, src_operator_scheme: str,
-                            known_hash: str | None,
-                            to_storage: OpendalClient, error_queue):
+    def _transfer_bg_thread(
+        self,
+        src_path: PathBuf,
+        src_operator: Operator,
+        src_operator_scheme: str,
+        known_hash: str | None,
+        to_storage: OpendalClient,
+        error_queue,
+    ):
         """Transfer image to storage in the background
 
         Args:
@@ -285,7 +289,6 @@ def _transfer_bg_thread(self, src_path: PathBuf, src_operator: Operator, src_ope
             raise
 
     def _sha256_file(self, src_operator, src_path) -> str:
-
         m = hashlib.sha256()
         with src_operator.open(src_path, "rb") as f:
             while True:
@@ -299,11 +302,13 @@ def _sha256_file(self, src_operator, src_path) -> str:
     def _create_metadata_and_json(self, src_operator, src_path) -> tuple[Metadata, str]:
         """Create a metadata json string from a metadata object"""
         metadata = src_operator.stat(src_path)
-        return metadata, json.dumps({
-            "path": str(src_path),
-            "content_length": metadata.content_length,
-            "etag": metadata.etag,
-        })
+        return metadata, json.dumps(
+            {
+                "path": str(src_path),
+                "content_length": metadata.content_length,
+                "etag": metadata.etag,
+            }
+        )
 
     def _lookup_block_device(self, console, prompt, address: str) -> str:
         """Lookup block device for a given address.
@@ -317,7 +322,7 @@ def _lookup_block_device(self, console, prompt, address: str) -> str:
         # lrwxrwxrwx    1 root     root             0 Jan  1
         # 00:00 mmcblk1 -> ../../devices/platform/bus@100000/4fb0000.mmc/mmc_host/mmc1/mmc1:aaaa/block/mmcblk1
         output = console.before.decode(errors="ignore")
-        match = re.search(r'\s(\w+)\s->', output)
+        match = re.search(r"\s(\w+)\s->", output)
         if match:
             return "/dev/" + match.group(1)
         else:
@@ -359,15 +364,13 @@ def _services_up(self):
             self.http.stop()
             self.tftp.stop()
 
-
     def _generate_uboot_env(self):
         """Generate a uboot environment dictionary, may need specific overrides for different targets"""
         tftp_host = self.tftp.get_host()
         return {
             "serverip": tftp_host,
         }
 
-
     @contextmanager
     def _busybox(self):
         """Start a busybox shell.
@@ -406,7 +409,7 @@ def _busybox(self):
                 uboot.run_command(f"tftpboot {dtb_address} {dtb_filename}", timeout=120)
 
             self.logger.info(f"Running boot command: {manifest.spec.bootcmd}")
-            console.send(manifest.spec.bootcmd +"\n")
+            console.send(manifest.spec.bootcmd + "\n")
 
             # if manifest has login details, we need to login
             if manifest.spec.login.username:
@@ -438,7 +441,7 @@ def use_initram(self, path: PathBuf, operator: Operator | None = None):
     def use_kernel(self, path: PathBuf, operator: Operator | None = None):
         """Use kernel file"""
         if operator is None:
-            path, operator, operator_scheme  = operator_for_path(path)
+            path, operator, operator_scheme = operator_for_path(path)
 
         ...
 
@@ -461,27 +464,37 @@ def base():
         @base.command()
         @click.argument("file")
         @click.option("--partition", type=str)
-        @click.option('--os-image-checksum',
-                        help='SHA256 checksum of OS image (direct value)')
-        @click.option('--os-image-checksum-file',
-                        help='File containing SHA256 checksum of OS image',
-                        type=click.Path(exists=True, dir_okay=False))
-        @click.option('--force-exporter-http', is_flag=True, help='Force use of exporter HTTP')
-        @click.option('--force-flash-bundle', type=str, help='Force use of a specific flasher OCI bundle')
+        @click.option("--os-image-checksum", help="SHA256 checksum of OS image (direct value)")
+        @click.option(
+            "--os-image-checksum-file",
+            help="File containing SHA256 checksum of OS image",
+            type=click.Path(exists=True, dir_okay=False),
+        )
+        @click.option("--force-exporter-http", is_flag=True, help="Force use of exporter HTTP")
+        @click.option("--force-flash-bundle", type=str, help="Force use of a specific flasher OCI bundle")
         @debug_console_option
-        def flash(file, partition, os_image_checksum, os_image_checksum_file,
-                  console_debug, force_exporter_http, force_flash_bundle):
+        def flash(
+            file,
+            partition,
+            os_image_checksum,
+            os_image_checksum_file,
+            console_debug,
+            force_exporter_http,
+            force_flash_bundle,
+        ):
             """Flash image to DUT from file"""
             if os_image_checksum_file and os.path.exists(os_image_checksum_file):
                 with open(os_image_checksum_file) as f:
                     os_image_checksum = f.read().strip().split()[0]
                     self.logger.info(f"Read checksum from file: {os_image_checksum}")
 
             self.set_console_debug(console_debug)
-            self.flash(file,
-                       partition=partition,
-                       force_exporter_http=force_exporter_http,
-                       force_flash_bundle=force_flash_bundle)
+            self.flash(
+                file,
+                partition=partition,
+                force_exporter_http=force_exporter_http,
+                force_flash_bundle=force_flash_bundle,
+            )
 
         @base.command()
         @debug_console_option
@@ -522,8 +535,8 @@ def _get_decompression_command(filename_or_url) -> str:
         filename = urlparse(filename_or_url).path.split("/")[-1]
 
     filename = filename.lower()
-    if filename.endswith(('.gz', '.gzip')):
-        return 'zcat |'
-    elif filename.endswith('.xz'):
-        return 'xzcat |'
-    return ''
+    if filename.endswith((".gz", ".gzip")):
+        return "zcat |"
+    elif filename.endswith(".xz"):
+        return "xzcat |"
+    return ""

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/driver.py

@@ -13,38 +13,41 @@
 
 @dataclass(kw_only=True)
 class BaseFlasher(Driver):
-    """ driver for Jumpstarter"""
+    """driver for Jumpstarter"""
+
     flasher_bundle: str = field(default="quay.io/jumpstarter-dev/jumpstarter-flasher-test:latest")
     cache_dir: str = field(default="/var/lib/jumpstarter/flasher")
-    tftp_dir : str = field(default="/var/lib/tftpboot")
-    http_dir : str = field(default="/var/www/html")
+    tftp_dir: str = field(default="/var/lib/tftpboot")
+    http_dir: str = field(default="/var/www/html")
 
     def __post_init__(self):
         if hasattr(super(), "__post_init__"):
             super().__post_init__()
 
         # Ensure required children are present if not already instantiated
         # in configuration
-        if 'tftp' not in self.children:
-            self.children['tftp'] = Tftp(root_dir=self.tftp_dir)
-            self.tftp = self.children['tftp']
+        if "tftp" not in self.children:
+            self.children["tftp"] = Tftp(root_dir=self.tftp_dir)
+            self.tftp = self.children["tftp"]
 
-        if 'http' not in self.children:
-            self.children['http'] = HttpServer(root_dir=self.http_dir)
-            self.http = self.children['http']
+        if "http" not in self.children:
+            self.children["http"] = HttpServer(root_dir=self.http_dir)
+            self.http = self.children["http"]
 
         # Ensure required children are present, the following are not auto-created
-        if 'serial' not in self.children:
-            raise ConfigurationError("'serial' instance is required for BaseFlasher "
-                                     "either via a ref ir a direct child instance")
+        if "serial" not in self.children:
+            raise ConfigurationError(
+                "'serial' instance is required for BaseFlasher either via a ref ir a direct child instance"
+            )
 
-        if 'power' not in self.children:
-            raise ConfigurationError("'power' instance is required for BaseFlasher "
-                                     "either via a ref ir a direct child instance")
+        if "power" not in self.children:
+            raise ConfigurationError(
+                "'power' instance is required for BaseFlasher either via a ref ir a direct child instance"
+            )
 
         # bundles that have already been downloaded in the current session
         self._downloaded = {}
-        self._use_dtb = None # use default dtb unless set by client
+        self._use_dtb = None  # use default dtb unless set by client
 
     @classmethod
     def client(cls) -> str:
@@ -76,7 +79,6 @@ async def setup_flasher_bundle(self, force_flash_bundle: str | None = None):
         self.logger.info(f"Setting up dtb in tftp: {dtb_path}")
         await self.tftp.storage.copy_exporter_file(dtb_path, dtb_path.name)
 
-
     @export
     def set_dtb(self, handle):
         """Provide a different dtb from client"""
@@ -87,8 +89,10 @@ async def use_dtb_variant(self, variant):
         """Provide a different dtb reference from the flasher bundle"""
         manifest = await self.get_flasher_manifest()
         if manifest.get_dtb_file(variant) is None:
-            raise ValueError(f"DTB variant {variant} not found in the flasher bundle, "
-                             f"available variants are: {list(manifest.spec.dtb.variants.keys())}")
+            raise ValueError(
+                f"DTB variant {variant} not found in the flasher bundle, "
+                f"available variants are: {list(manifest.spec.dtb.variants.keys())}"
+            )
         self._use_dtb = variant
 
     def set_kernel(self, handle):
@@ -186,5 +190,6 @@ async def get_initram_address(self) -> str:
 
 @dataclass(kw_only=True)
 class TIJ784S4Flasher(BaseFlasher):
-    """ driver for Jumpstarter"""
+    """driver for Jumpstarter"""
+
     flasher_bundle: str = "quay.io/jumpstarter-dev/jumpstarter-flasher-ti-j784s4:latest"