From 826bb4f2ddaa6e6d1401553be3f8b3fe9cba4357 Mon Sep 17 00:00:00 2001
From: Newbyte <newbyte@postmarketos.org>
Date: Tue, 29 Oct 2024 17:27:50 +0100
Subject: [PATCH] pmb: Properly type Bootimg (MR 2464)

This could be done better with a real class instead of a TypedDict, but
it's better than a regular dict.
---
 pmb/aportgen/device.py | 32 ++++++++++++++++++--------------
 pmb/parse/bootimg.py   | 20 ++++++++++++++++++--
 pmb/types.py           | 17 +++++++++++++++++
 3 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/pmb/aportgen/device.py b/pmb/aportgen/device.py
index 0b51ff2e..268de250 100644
--- a/pmb/aportgen/device.py
+++ b/pmb/aportgen/device.py
@@ -3,6 +3,7 @@
 from pmb.core.context import get_context
 from pmb.core.arch import Arch
 from pmb.helpers import logging
+from pmb.types import Bootimg
 from pathlib import Path
 import os
 import pmb.helpers.cli
@@ -119,21 +120,24 @@ def ask_for_bootimg():
             logging.fatal("ERROR: " + str(e) + ". Please try again.")
 
 
-def generate_deviceinfo_fastboot_content(bootimg=None):
+def generate_deviceinfo_fastboot_content(bootimg: Bootimg | None = None) -> str:
     if bootimg is None:
-        bootimg = {
-            "cmdline": "",
-            "qcdt": "false",
-            "dtb_second": "false",
-            "base": "",
-            "kernel_offset": "",
-            "ramdisk_offset": "",
-            "second_offset": "",
-            "tags_offset": "",
-            "pagesize": "2048",
-            "mtk_label_kernel": "",
-            "mtk_label_ramdisk": "",
-        }
+        bootimg = Bootimg(
+            cmdline="",
+            qcdt="false",
+            qcdt_type=None,
+            dtb_offset=None,
+            dtb_second="false",
+            base="",
+            kernel_offset="",
+            ramdisk_offset="",
+            second_offset="",
+            tags_offset="",
+            pagesize="2048",
+            header_version=None,
+            mtk_label_kernel="",
+            mtk_label_ramdisk="",
+        )
 
     content = f"""\
         deviceinfo_kernel_cmdline="{bootimg["cmdline"]}"
diff --git a/pmb/parse/bootimg.py b/pmb/parse/bootimg.py
index edc026ae..342016d7 100644
--- a/pmb/parse/bootimg.py
+++ b/pmb/parse/bootimg.py
@@ -9,6 +9,7 @@ import pmb.chroot
 import pmb.chroot.other
 import pmb.chroot.apk
 from pmb.core import Chroot
+from pmb.types import Bootimg
 
 
 def is_dtb(path) -> bool:
@@ -73,7 +74,7 @@ def get_qcdt_type(path) -> str | None:
             return None
 
 
-def bootimg(path: Path) -> dict[str, str]:
+def bootimg(path: Path) -> Bootimg:
     if not path.exists():
         raise RuntimeError(f"Could not find file '{path}'")
 
@@ -176,7 +177,22 @@ def bootimg(path: Path) -> dict[str, str]:
     # Cleanup
     pmb.chroot.user(["rm", "-r", temp_path])
 
-    return output
+    return Bootimg(
+        cmdline=output["cmdline"],
+        qcdt=output["qcdt"],
+        qcdt_type=output.get("qcdt_type"),
+        dtb_offset=output.get("dtb_offset"),
+        dtb_second=output["dtb_second"],
+        base=output["base"],
+        kernel_offset=output["kernel_offset"],
+        ramdisk_offset=output["ramdisk_offset"],
+        second_offset=output["second_offset"],
+        tags_offset=output["tags_offset"],
+        pagesize=output["pagesize"],
+        header_version=output.get("header_version"),
+        mtk_label_kernel=output["mtk_label_kernel"],
+        mtk_label_ramdisk=output["mtk_label_ramdisk"],
+    )
 
 
 def trim_input(f) -> str:
diff --git a/pmb/types.py b/pmb/types.py
index e7aea670..0d283eea 100644
--- a/pmb/types.py
+++ b/pmb/types.py
@@ -29,6 +29,23 @@ class AportGenEntry(TypedDict):
     confirm_overwrite: bool
 
 
+class Bootimg(TypedDict):
+    cmdline: str
+    qcdt: str
+    qcdt_type: str | None
+    dtb_offset: str | None
+    dtb_second: str
+    base: str
+    kernel_offset: str
+    ramdisk_offset: str
+    second_offset: str
+    tags_offset: str
+    pagesize: str
+    header_version: str | None
+    mtk_label_kernel: str
+    mtk_label_ramdisk: str
+
+
 # Property list generated with:
 # $ rg --vimgrep "((^|\s)args\.\w+)" --only-matching | cut -d"." -f3 | sort | uniq
 class PmbArgs(Namespace):