feat: add type hints to ubireader.ubi

Author: scootermon

ubireader/ubi/__init__.py

@@ -17,12 +17,19 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING
 from ubireader.debug import error
 from ubireader.ubi.block import sort, extract_blocks
 from ubireader.ubi import display
 from ubireader.ubi.image import description as image
 from ubireader.ubi.block import layout, rm_old_blocks
 
+if TYPE_CHECKING:
+    from ubireader.ubi_io import ubi_file as UbiFile
+    from ubireader.ubi.block import description as Block
+    from ubireader.ubi.image import description as Image
+
 class ubi_base(object):
     """UBI Base object
 
@@ -39,7 +46,7 @@ class ubi_base(object):
     Dict:blocks        -- Dict keyed by PEB number of all blocks.
     """
 
-    def __init__(self, ubi_file):
+    def __init__(self, ubi_file: UbiFile) -> None:
         self.__name__ = 'UBI'
         self._file = ubi_file
         self._first_peb_num = 0
@@ -54,7 +61,7 @@ def __init__(self, ubi_file):
         self._leb_size = self.file.block_size - arbitrary_block.ec_hdr.data_offset
 
 
-    def _get_file(self):
+    def _get_file(self) -> UbiFile:
         """UBI File object
 
         Returns:
@@ -64,7 +71,7 @@ def _get_file(self):
     file = property(_get_file)
 
 
-    def _get_block_count(self):
+    def _get_block_count(self) -> int:
         """Total amount of UBI blocks in file.
 
         Returns:
@@ -74,9 +81,9 @@ def _get_block_count(self):
     block_count = property(_get_block_count)
 
 
-    def _set_first_peb_num(self, i):
+    def _set_first_peb_num(self, i: int) -> None:
         self._first_peb_num = i
-    def _get_first_peb_num(self):
+    def _get_first_peb_num(self) -> int:
         """First Physical Erase Block with UBI data
         
         Returns:
@@ -86,7 +93,7 @@ def _get_first_peb_num(self):
     first_peb_num = property(_get_first_peb_num, _set_first_peb_num)
 
 
-    def _get_leb_size(self):
+    def _get_leb_size(self) -> int:
         """LEB size of UBI blocks in file.
 
         Returns:
@@ -96,7 +103,7 @@ def _get_leb_size(self):
     leb_size = property(_get_leb_size)
 
 
-    def _get_peb_size(self):
+    def _get_peb_size(self) -> int:
         """PEB size of UBI blocks in file.
 
         Returns:
@@ -106,7 +113,7 @@ def _get_peb_size(self):
     peb_size = property(_get_peb_size)
 
 
-    def _get_min_io_size(self):
+    def _get_min_io_size(self) -> int:
         """Min I/O Size
 
         Returns:
@@ -116,7 +123,7 @@ def _get_min_io_size(self):
     min_io_size = property(_get_min_io_size)
 
 
-    def _get_blocks(self):
+    def _get_blocks(self) -> dict[int, Block]:
         """Main Dict of UBI Blocks
 
         Passed around for lists of indexes to be made or to be returned
@@ -142,7 +149,7 @@ class ubi(ubi_base):
     List:unknown_blocks_list  -- List of blocks with unknown types. *
     """
 
-    def __init__(self, ubi_file):
+    def __init__(self, ubi_file: UbiFile) -> None:
         super(ubi, self).__init__(ubi_file)
 
         layout_list, data_list, int_vol_list, unknown_list = sort.by_type(self.blocks)
@@ -161,12 +168,12 @@ def __init__(self, ubi_file):
 
         layout_infos = layout.associate_blocks(self.blocks, layout_pairs)
 
-        self._images = []
+        self._images: list[Image] = []
         for i in range(0, len(layout_infos)):
             self._images.append(image(self.blocks, layout_infos[i]))
 
 
-    def _get_images(self):
+    def _get_images(self) -> list[Image]:
         """Get UBI images.
 
         Returns:
@@ -176,7 +183,7 @@ def _get_images(self):
     images = property(_get_images)
 
 
-    def _get_data_blocks_list(self):
+    def _get_data_blocks_list(self) -> list[int]:
         """Get all UBI blocks found in file that are data blocks.
 
         Returns:
@@ -186,7 +193,7 @@ def _get_data_blocks_list(self):
     data_blocks_list = property(_get_data_blocks_list)
 
 
-    def _get_layout_blocks_list(self):
+    def _get_layout_blocks_list(self) -> list[int]:
         """Get all UBI blocks found in file that are layout volume blocks.
 
         Returns:
@@ -196,7 +203,7 @@ def _get_layout_blocks_list(self):
     layout_blocks_list = property(_get_layout_blocks_list)
 
 
-    def _get_int_vol_blocks_list(self):
+    def _get_int_vol_blocks_list(self) -> list[int]:
         """Get all UBI blocks found in file that are internal volume blocks.
 
         Returns:
@@ -208,7 +215,7 @@ def _get_int_vol_blocks_list(self):
     int_vol_blocks_list = property(_get_int_vol_blocks_list)
 
 
-    def _get_unknown_blocks_list(self):
+    def _get_unknown_blocks_list(self) -> list[int]:
         """Get all UBI blocks found in file of unknown type..
 
         Returns:
@@ -217,7 +224,7 @@ def _get_unknown_blocks_list(self):
         return self._unknown_blocks_list
     unknown_blocks_list = property(_get_unknown_blocks_list)
 
-    def display(self, tab=''):
+    def display(self, tab: str = '') -> str:
         """Print information about this object.
         
         Argument:

ubireader/ubi/block/__init__.py

@@ -17,13 +17,19 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING
 from zlib import crc32
 from ubireader import settings
 from ubireader.debug import error, log, verbose_display, verbose_log
 from ubireader.ubi import display
 from ubireader.ubi.defines import UBI_EC_HDR_SZ, UBI_VID_HDR_SZ, UBI_INTERNAL_VOL_START, UBI_EC_HDR_MAGIC, UBI_CRC32_INIT
 from ubireader.ubi.headers import ec_hdr, vid_hdr, vtbl_recs
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+    from ubireader.ubi import ubi_base as UbiBase
+    from ubireader.ubi.headers import _vtbl_rec as VtblRec
 
 class description(object):
     """UBI Block description Object
@@ -48,16 +54,18 @@ class description(object):
     Will print out all information when invoked as a string.
     """
 
-    def __init__(self, block_buf):
+    data_crc: int
+
+    def __init__(self, block_buf: bytes) -> None:
 
         self.file_offset = -1
         self.peb_num = -1
         self.leb_num = -1
         self.size = -1
 
-        self.vid_hdr = None
+        self.vid_hdr: vid_hdr | None = None
         self.is_internal_vol = False
-        self.vtbl_recs = []
+        self.vtbl_recs: list[VtblRec] = []
 
         # TODO better understanding of block types/errors
         self.ec_hdr = ec_hdr(block_buf[0:UBI_EC_HDR_SZ])
@@ -77,16 +85,16 @@ def __init__(self, block_buf):
         self.is_valid = not self.ec_hdr.errors and not self.vid_hdr.errors or settings.ignore_block_header_errors
 
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return 'Block: PEB# %s: LEB# %s' % (self.peb_num, self.leb_num)
 
 
-    def display(self, tab=''):
+    def display(self, tab: str ='') -> str:
         return display.block(self, tab)
 
 
 
-def get_blocks_in_list(blocks, idx_list):
+def get_blocks_in_list(blocks: Mapping[int, description], idx_list: Iterable[int]) -> dict[int, description]:
     """Retrieve block objects in list of indexes
 
     Arguments:
@@ -103,7 +111,7 @@ def get_blocks_in_list(blocks, idx_list):
 
 
 
-def extract_blocks(ubi):
+def extract_blocks(ubi: UbiBase) -> dict[int, description]:
     """Get a list of UBI block objects from file
 
     Arguments:.
@@ -113,11 +121,11 @@ def extract_blocks(ubi):
     Dict -- Of block objects keyed by PEB number.
     """
 
-    blocks = {}
+    blocks: dict[int, description] = {}
     ubi.file.seek(ubi.file.start_offset)
     peb_count = 0
     cur_offset = 0
-    bad_blocks = []
+    bad_blocks: list[int] = []
 
     # range instead of xrange, as xrange breaks > 4GB end_offset.
     for i in range(ubi.file.start_offset, ubi.file.end_offset, ubi.file.block_size):
@@ -157,8 +165,8 @@ def extract_blocks(ubi):
     return blocks
 
 
-def rm_old_blocks(blocks, block_list):
-    del_blocks = []
+def rm_old_blocks(blocks: Mapping[int, description], block_list: Iterable[int]) -> list[int]:
+    del_blocks: list[int] = []
 
     for i in block_list:
         if i in del_blocks:

ubireader/ubi/block/layout.py

@@ -17,10 +17,20 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING, Literal, Protocol, cast, overload
 from ubireader.debug import log
 from ubireader.ubi.block import sort
 
-def group_pairs(blocks, layout_blocks_list):
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+    from ubireader.ubi.block import description as Block
+
+class _LayoutPair(Protocol):
+    def __getitem__(self, idx: Literal[0, 1]) -> int: ...
+    def __contains__(self, item: int) -> bool: ...
+
+def group_pairs(blocks: Mapping[int, Block], layout_blocks_list: Iterable[int]) -> list[_LayoutPair]:
     """Sort a list of layout blocks into pairs
 
     Arguments:
@@ -31,7 +41,7 @@ def group_pairs(blocks, layout_blocks_list):
     List -- Layout block pair indexes grouped in a list
     """
 
-    image_dict={}
+    image_dict: dict[int, _LayoutPair] = {}
     for block_id in layout_blocks_list:
         image_seq=blocks[block_id].ec_hdr.image_seq
         if image_seq not in image_dict:
@@ -43,8 +53,13 @@ def group_pairs(blocks, layout_blocks_list):
 
     return list(image_dict.values())
 
+class _LayoutInfo(_LayoutPair):
+    @overload
+    def __getitem__(self, idx: Literal[0, 1]) -> int: ...
+    @overload
+    def __getitem__(self, idx: Literal[2]) -> list[int]: ...
 
-def associate_blocks(blocks, layout_pairs):
+def associate_blocks(blocks: Mapping[int, Block], layout_pairs: list[_LayoutPair]) -> list[_LayoutInfo]:
     """Group block indexes with appropriate layout pairs
 
     Arguments:
@@ -55,9 +70,9 @@ def associate_blocks(blocks, layout_pairs):
     List -- Layout block pairs grouped with associated block ranges.
     """
 
-    seq_blocks = []
+    seq_blocks: list[int] = []
     for layout_pair in layout_pairs:
-        seq_blocks = sort.by_image_seq(blocks, blocks[layout_pair[0]].ec_hdr.image_seq)
+        seq_blocks = sort.by_image_seq(blocks, blocks[cast(int, layout_pair[0])].ec_hdr.image_seq)
         seq_blocks = [b for b in seq_blocks if b not in layout_pair]
         layout_pair.append(seq_blocks)
 

ubireader/ubi/block/sort.py

@@ -17,9 +17,15 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING, Literal
 from ubireader import settings
 
-def by_image_seq(blocks, image_seq):
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from ubireader.ubi.block import description as Block
+
+def by_image_seq(blocks: Mapping[int, Block], image_seq: int) -> list[int]:
     """Filter blocks to return only those associated with the provided image_seq number.
        If uboot_fix is set, associate blocks with an image_seq of 0 also.
 
@@ -36,7 +42,7 @@ def by_image_seq(blocks, image_seq):
     else:
         return list(filter(lambda block: blocks[block].ec_hdr.image_seq == image_seq, blocks))
 
-def by_leb(blocks):
+def by_leb(blocks: Mapping[int, Block]) -> list[Literal['x'] | int]:
     """Sort blocks by Logical Erase Block number.
     
     Arguments:
@@ -46,7 +52,7 @@ def by_leb(blocks):
     List              -- Indexes of blocks sorted by LEB.
     """ 
     slist_len = len(blocks)
-    slist = ['x'] * slist_len
+    slist: list[Literal['x'] | int] = ['x'] * slist_len
 
     for block in blocks:
         if blocks[block].leb_num >= slist_len:
@@ -59,7 +65,7 @@ def by_leb(blocks):
     return slist
 
 
-def by_vol_id(blocks, slist=None):
+def by_vol_id(blocks: Mapping[int, Block], slist: list[int] | None = None) -> dict[int, list[int]]:
     """Sort blocks by volume id
 
     Arguments:
@@ -70,7 +76,7 @@ def by_vol_id(blocks, slist=None):
     Dict -- blocks grouped in lists with dict key as volume id.
     """
 
-    vol_blocks = {}
+    vol_blocks: dict[int, list[int]] = {}
 
     # sort block by volume
     # not reliable with multiple partitions (fifo)
@@ -88,7 +94,7 @@ def by_vol_id(blocks, slist=None):
 
     return vol_blocks
 
-def by_type(blocks, slist=None):
+def by_type(blocks: Mapping[int, Block], slist: list[int] | None = None) -> tuple[list[int], list[int], list[int], list[int]]:
     """Sort blocks into layout, internal volume, data or unknown
 
     Arguments:
@@ -106,10 +112,10 @@ def by_type(blocks, slist=None):
                     of crc in ed_hdr or vid_hdr.
     """
 
-    layout = []
-    data = []
-    int_vol = []
-    unknown = []
+    layout: list[int] = []
+    data: list[int] = []
+    int_vol: list[int] = []
+    unknown: list[int] = []
 
     for i in blocks:
         if slist and i not in slist: