Decouple input parsing from XML generation

.github/workflows/build.yml

@@ -14,11 +14,11 @@ jobs:
         with:
           fetch-depth: 0
 
-      - name: Install linters
+      - name: Install linters and test runner
         run: |
           sudo snap install ruff
           sudo apt-get update
-          sudo apt-get install -y mypy
+          sudo apt-get install -y mypy python3-pytest
 
       - name: Install qcom-ptool
         run: |
@@ -29,6 +29,7 @@ jobs:
           PTOOL_SEED: qcom-ptool-ci
         run: |
           make lint
+          make unit-test
           make all integration check-checksums
 
       - name: Verify checksum manifest is up to date

Makefile

@@ -12,7 +12,7 @@ QCOM_PTOOL ?= qcom-ptool
 # optional build_id for Axiom contents.xml files
 BUILD_ID ?=
 
-.PHONY: all check check-checksums clean generate-checksums install lint integration
+.PHONY: all check check-checksums clean generate-checksums install lint integration unit-test
 
 all: $(PLATFORMS) $(PARTITIONS_XML) $(CONTENTS_XML)
 
@@ -29,6 +29,9 @@ lint:
 	ruff check qcom_ptool
 	mypy qcom_ptool
 
+unit-test:
+	pytest
+
 integration: all
 	# make sure generated output has created expected files
 	tests/integration/check-missing-files platforms/*/*/*.xml
@@ -46,7 +49,7 @@ generate-checksums: all
 	  ! -name '*.xml.in' -print0 | LC_ALL=C sort -z | xargs -0 sha256sum \
 	  > tests/integration/checksums.sha256
 
-check: lint integration
+check: lint unit-test integration
 
 install:
 	pip install .

README.md

@@ -33,14 +33,14 @@ subcommand.
 At runtime, the scripts use only the Python standard library (Python 3.8+),
 so no runtime dependencies need to be installed beyond the package itself.
 
-For development, `make lint` invokes `ruff` and `mypy` directly from the
-command line. On Debian/Ubuntu, install them as follows (ruff is not
-packaged in apt on all releases/architectures, so we install it from
-snap):
+For development, `make lint` invokes `ruff` and `mypy` and `make unit-test`
+runs the `pytest` suite under `tests/unit/`. On Debian/Ubuntu, install
+them as follows (ruff is not packaged in apt on all releases/architectures,
+so we install it from snap):
 
 ```sh
 sudo snap install ruff
-sudo apt install mypy
+sudo apt install mypy python3-pytest
 ```
 
 ## Makefile targets
@@ -49,8 +49,9 @@ sudo apt install mypy
 |---------------|------------------------------------------------------------|
 | `all`         | Generate partition XML and GPT binaries for all platforms  |
 | `lint`        | Run ruff (linter) and mypy (type checker) on the package   |
+| `unit-test`   | Run the pytest suite under `tests/unit/`                   |
 | `integration` | Build all platforms and verify generated files are present |
-| `check`       | Run both `lint` and `integration`                          |
+| `check`       | Run `lint`, `unit-test`, and `integration`                 |
 | `install`     | Install the package (`pip install .`)                      |
 | `clean`       | Remove generated XML and binary files from platforms/      |
 
@@ -63,12 +64,13 @@ The Makefile invokes `qcom-ptool` from `PATH`. Install the package (or
 # install the tool
 pip install -e .
 
-# install linters (Debian/Ubuntu)
+# install linters and test runner (Debian/Ubuntu)
 sudo snap install ruff
-sudo apt install mypy
+sudo apt install mypy python3-pytest
 
-# run linters
+# run linters and unit tests
 make lint
+make unit-test
 
 # build all platforms and run tests
 make check

pyproject.toml

@@ -53,3 +53,10 @@ ignore_missing_imports = true
 module = ["qcom_ptool.msp", "qcom_ptool.ptool"]
 # Legacy scripts - enable gradually
 check_untyped_defs = false
+
+[tool.pytest.ini_options]
+# Allow running `pytest` from the repo root without requiring an editable
+# install of the package; tests under tests/unit/ import qcom_ptool directly
+# from the source tree.
+pythonpath = ["."]
+testpaths = ["tests/unit"]

qcom_ptool/gen_partition.py

@@ -27,14 +27,17 @@
 # OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
 # IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+from __future__ import annotations
+
 import getopt
-import re
 import sys
 import xml.etree.ElementTree as ET
-from collections import OrderedDict
 from typing import NoReturn
 from xml.dom import minidom
 
+from qcom_ptool.loaders import load as load_spec
+from qcom_ptool.spec import DiskParams, PartitionsByLun
+
 
 def usage() -> NoReturn:
     print(
@@ -44,159 +47,9 @@ def usage() -> NoReturn:
     sys.exit(1)
 
 
-##################################################################
-# defaults to be used
-disk_params_defaults = OrderedDict(
-    {
-        "type": "",
-        "size": "",
-        "SECTOR_SIZE_IN_BYTES": "512",
-        "WRITE_PROTECT_BOUNDARY_IN_KB": "65536",
-        "GROW_LAST_PARTITION_TO_FILL_DISK": "false",
-        "ALIGN_PARTITIONS_TO_PERFORMANCE_BOUNDARY": "true",
-        "PERFORMANCE_BOUNDARY_IN_KB": "4",
-    }
-)
-
-partition_entry_defaults = {
-    "label": "",
-    "size_in_kb": "",
-    "type": "00000000-0000-0000-0000-000000000000",
-    "bootable": "false",
-    "readonly": "true",
-    "filename": "",
-    "sparse": "false",
-}
-
-##################################################################
-# store entries read from input file
-disk_entry = None
-partition_entries = []
-# store partition image map passed from command line
-partition_image_map: dict[str, str] = {}
-input_file = None
-output_xml = None
-
-
-def disk_options(argv):
-    disk_params = disk_params_defaults.copy()
-    for opt, arg in argv:
-        if opt in ["--type"]:
-            disk_params["type"] = arg
-        elif opt in ["--size"]:
-            disk_params["size"] = arg
-        elif opt in ["--sector-size-in-bytes"]:
-            disk_params["SECTOR_SIZE_IN_BYTES"] = arg
-        elif opt in ["--write-protect-boundary"]:
-            disk_params["WRITE_PROTECT_BOUNDARY_IN_KB"] = arg
-        elif opt in ["--grow-last-partition"]:
-            disk_params["GROW_LAST_PARTITION_TO_FILL_DISK"] = "true"
-        elif opt in ["--align-partitions"]:
-            disk_params["ALIGN_PARTITIONS_TO_PERFORMANCE_BOUNDARY"] = "true"
-            disk_params["PERFORMANCE_BOUNDARY_IN_KB"] = str(int(arg) // 1024)
-    return disk_params
-
-
-def partition_size_in_kb(size):
-    if not re.search("[a-zA-Z]+", size):
-        return int(size) // 1024
-    m = re.search("([0-9]+)(?=[Kk][Bb]?)", size)
-    if m:
-        return int(m.group(0))
-    m = re.search("([0-9]+)(?=[Mm][Bb]?)", size)
-    if m:
-        return int(m.group(0)) * 1024
-    m = re.search("([0-9]+)(?=[Gg][Bb]?)", size)
-    if m:
-        return int(m.group(0)) * 1024 * 1024
-    raise ValueError("Unrecognized size format: '%s'" % size)
-
-
-def partition_options(argv):
-    partition_entry = partition_entry_defaults.copy()
-    phys_part = 0
-    for opt, arg in argv:
-        if opt in ["--lun", "--phys-part"]:
-            phys_part = arg
-        elif opt in ["--name"]:
-            partition_entry["label"] = arg
-        elif opt in ["--size"]:
-            kbytes = partition_size_in_kb(arg)
-            partition_entry["size_in_kb"] = str(kbytes)
-        elif opt in ["--type-guid"]:
-            partition_entry["type"] = arg
-        elif opt in ["--attributes"]:
-            attribute_bits = int(arg, 16)
-            if attribute_bits & (1 << 2):
-                partition_entry["bootable"] = "true"
-            else:
-                partition_entry["bootable"] = "false"
-            if attribute_bits & (1 << 60):
-                partition_entry["readonly"] = "true"
-            else:
-                partition_entry["readonly"] = "false"
-        elif opt in ["--filename"]:
-            partition_entry["filename"] = arg
-        elif opt in ["--sparse"]:
-            partition_entry["sparse"] = arg
-        if partition_entry["label"] in partition_image_map:
-            partition_entry["filename"] = partition_image_map[partition_entry["label"]]
-    return phys_part, partition_entry
-
-
-def parse_partition_entries(partition_entries):
-    partitions_params: dict[int, list[dict[str, str]]] = {}
-
-    for partition_entry in partition_entries:
-        opts_list = list(partition_entry.split(" "))
-        if opts_list[0] == "--partition":
-            try:
-                options, _remainders = getopt.gnu_getopt(
-                    opts_list[1:],
-                    "",
-                    [
-                        "lun=",
-                        "phys-part=",
-                        "name=",
-                        "size=",
-                        "type-guid=",
-                        "filename=",
-                        "attributes=",
-                        "sparse=",
-                    ],
-                )
-                phys_part, partition = partition_options(options)
-                partitions_params.setdefault(phys_part, []).append(partition)
-            except Exception as e:
-                print(str(e))
-                usage()
-
-    return partitions_params
-
-
-def parse_disk_entry(disk_entry):
-    opts_list = list(disk_entry.split(" "))
-    if opts_list[0] == "--disk":
-        try:
-            options, _remainders = getopt.gnu_getopt(
-                opts_list[1:],
-                "",
-                [
-                    "type=",
-                    "size=",
-                    "sector-size-in-bytes=",
-                    "write-protect-boundary=",
-                    "grow-last-partition",
-                    "align-partitions=",
-                ],
-            )
-            return disk_options(options)
-        except Exception as e:
-            print(str(e))
-            usage()
-
-
-def generate_multi_lun_xml(disk_params, partitions, output_xml):
+def generate_multi_lun_xml(
+    disk_params: DiskParams, partitions: PartitionsByLun, output_xml: str
+) -> None:
     root = ET.Element("configuration")
     parser_instruction_text = ""
 
@@ -222,7 +75,9 @@ def generate_multi_lun_xml(disk_params, partitions, output_xml):
         f.write(xmlstr)
 
 
-def generate_partition_xml(disk_params, partitions, output_xml):
+def generate_partition_xml(
+    disk_params: DiskParams, partitions: PartitionsByLun, output_xml: str
+) -> None:
     print("Generating %s XML %s" % (disk_params["type"].upper(), output_xml))
 
     if disk_params["type"] in ("emmc", "nvme", "spinor", "ufs"):
@@ -236,58 +91,51 @@ def generate_partition_xml(disk_params, partitions, output_xml):
 
 ###############################################################################
 # main
-disk_entry_err_msg = "contains more than one --disk entries"
 
-if len(sys.argv) < 3:
-    usage()
 
-try:
-    if sys.argv[1] == "-h" or sys.argv[1] == "--help":
+def main(argv: list[str] | None = None) -> int:
+    if argv is None:
+        argv = sys.argv
+
+    if len(argv) < 3:
+        usage()
+
+    input_file: str | None = None
+    output_xml: str | None = None
+    image_map: dict[str, str] = {}
+
+    if argv[1] == "-h" or argv[1] == "--help":
         usage()
     try:
-        opts, rem = getopt.getopt(sys.argv[1:], "i:o:m:")
+        opts, _rem = getopt.getopt(argv[1:], "i:o:m:")
         for opt, arg in opts:
-            if opt in ["-i"]:
+            if opt == "-i":
                 input_file = arg
-            elif opt in ["-o"]:
+            elif opt == "-o":
                 output_xml = arg
-            elif opt in ["-m"]:
+            elif opt == "-m":
                 for mapping in arg.split(","):
                     tags = mapping.split("=")
                     if len(tags) > 1:
-                        partition_image_map[tags[0]] = tags[1]
+                        image_map[tags[0]] = tags[1]
             else:
                 usage()
-
     except Exception as argerr:
         print(str(argerr))
         usage()
+
     if input_file is None or output_xml is None:
         usage()
-    f = open(input_file)
-    line = f.readline()
-    while line:
-        if not re.search(r"^\s*#", line) and not re.search(r"^\s*$", line):
-            line = line.strip()
-            if re.search("^--disk", line):
-                if disk_entry is None:
-                    disk_entry = line
-                else:
-                    print("%s %s" % (sys.argv[1], disk_entry_err_msg))
-                    print("%s\n%s" % (disk_entry, line))
-                    sys.exit(1)
-            elif re.search("^--partition", line):
-                partition_entries.append(line)
-            else:
-                print("Ignoring %s" % (line))
-        line = f.readline()
-    f.close()
-except Exception as e:
-    print("Error: ", e)
-    sys.exit(1)
 
-disk_params = parse_disk_entry(disk_entry)
-partitions = parse_partition_entries(partition_entries)
-generate_partition_xml(disk_params, partitions, output_xml)
+    try:
+        spec = load_spec(input_file, image_map=image_map)
+    except Exception as e:
+        print("Error: ", e)
+        return 1
+
+    generate_partition_xml(spec["disk"], spec["partitions"], output_xml)
+    return 0
+
 
-sys.exit(0)
+if __name__ == "__main__":
+    sys.exit(main())