diff --git a/src/smpclient/mcuboot.py b/src/smpclient/mcuboot.py index ebf17fc..78aa33b 100644 --- a/src/smpclient/mcuboot.py +++ b/src/smpclient/mcuboot.py @@ -8,15 +8,13 @@ import argparse import pathlib import struct +from dataclasses import dataclass from enum import IntEnum, IntFlag, unique from functools import cached_property from io import BufferedReader, BytesIO -from typing import Annotated, Any, Final, Generic, Literal, TypeVar, Union +from typing import Final, Generic, Literal, TypeVar from intelhex import hex2bin # type: ignore -from pydantic import Field, GetCoreSchemaHandler -from pydantic.dataclasses import dataclass -from pydantic_core import CoreSchema, core_schema ImageMagic = Literal[0x96F3B83D] IMAGE_MAGIC: Final[ImageMagic] = 0x96F3B83D @@ -146,31 +144,25 @@ def __new__(cls, value: int) -> 'VendorTLV': ) return int.__new__(cls, value) - @classmethod - def __get_pydantic_core_schema__( - cls, _source_type: Any, _handler: GetCoreSchemaHandler - ) -> CoreSchema: - def validate(value: int) -> VendorTLV: - return cls(value) - return core_schema.no_info_after_validator_function( - validate, - core_schema.int_schema(), - ) - - -ImageTLVType = Annotated[Union[IMAGE_TLV, VendorTLV, int], Field(union_mode="left_to_right")] +ImageTLVType = IMAGE_TLV | VendorTLV | int """TLV type that accepts standard IMAGE_TLV enums, vendor-defined TLVs, or any integer. -This uses Pydantic's "left to right" union mode to: -1. First try to match against IMAGE_TLV enum values -2. Then try to validate as a VendorTLV (0xXXA0-0xXXFE ranges) -3. Finally accept any integer as a fallback - -This ensures backward compatibility and supports future TLV types without validation errors. +`ImageTLV` narrows a raw type field to the leftmost member that accepts it, so an +unrecognized type stays readable as an `int` instead of failing the parse. """ +def _narrow_tlv_type(value: int) -> ImageTLVType: + """Return `value` as the leftmost `ImageTLVType` that accepts it.""" + for tlv_type in (IMAGE_TLV, VendorTLV): + try: + return tlv_type(value) + except ValueError: + continue + return value + + @dataclass(frozen=True) class ImageVersion: """An MCUBoot image_version struct.""" @@ -223,7 +215,7 @@ def loads(data: bytes) -> ImageHeader: hdr_size=hdr_size, protect_tlv_size=protect_tlv_size, img_size=img_size, - flags=flags, + flags=IMAGE_F(flags), ver=ImageVersion(*ver), ) @@ -282,6 +274,10 @@ class ImageTLV: len: int """Data length (not including TLV header).""" + def __post_init__(self) -> None: + """Narrow `type` to the leftmost `ImageTLVType` that accepts it.""" + object.__setattr__(self, "type", _narrow_tlv_type(self.type)) + @staticmethod def load_from(file: BytesIO | BufferedReader) -> ImageTLV: """Load an `ImageTLV` from a file.""" diff --git a/tests/test_mcuboot_tools.py b/tests/test_mcuboot_tools.py index 3b2b8a8..2fbddeb 100644 --- a/tests/test_mcuboot_tools.py +++ b/tests/test_mcuboot_tools.py @@ -21,12 +21,12 @@ ImageTLVInfo, ImageTLVInfoMagic, ImageTLVProtInfoMagic, - ImageTLVType, ImageTLVValue, ImageVersion, MCUBootImageError, TLVNotFound, VendorTLV, + _narrow_tlv_type, mcuimg, ) @@ -195,24 +195,20 @@ def test_unknown_tlv_fallback() -> None: def test_tlv_type_union_order() -> None: """Test that union resolution follows left-to-right order.""" - from pydantic import TypeAdapter - - adapter: TypeAdapter[ImageTLVType] = TypeAdapter(ImageTLVType) - # Standard TLV should match IMAGE_TLV first - result = adapter.validate_python(0x02) - assert isinstance(result, IMAGE_TLV) - assert result == IMAGE_TLV.PUBKEY - - # Vendor TLV should validate - result = adapter.validate_python(0xA0) - assert isinstance(result, int) - assert result == 0xA0 - - # Unknown TLV should fallback to int - result = adapter.validate_python(0x99) - assert isinstance(result, int) - assert result == 0x99 + standard = _narrow_tlv_type(0x02) + assert isinstance(standard, IMAGE_TLV) + assert standard == IMAGE_TLV.PUBKEY + + # Vendor TLV should validate; VendorTLV rather than a bare int is the whole point + vendor = _narrow_tlv_type(0xA0) + assert isinstance(vendor, VendorTLV) + assert vendor == 0xA0 + + # Unknown TLV should fallback to int, matching neither of the narrower members + unknown = _narrow_tlv_type(0x99) + assert not isinstance(unknown, (IMAGE_TLV, VendorTLV)) + assert unknown == 0x99 def test_tlv_value_str_standard() -> None: