feat: add type hints to top-level ubireader module

Author: scootermon

ubireader/debug.py

@@ -17,23 +17,36 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
 import sys
 import traceback
+from typing import Literal, NoReturn, Protocol, overload
 from ubireader import settings
 
-def log(obj, message):
+class _Obj(Protocol):
+    __name__: str
+
+def log(obj: _Obj, message: str) -> None:
     if settings.logging_on or settings.logging_on_verbose:
         print('{} {}'.format(obj.__name__, message))
 
-def verbose_log(obj, message):
+def verbose_log(obj: _Obj, message: str) -> None:
     if settings.logging_on_verbose:
         log(obj, message)
 
-def verbose_display(displayable_obj):
+class _Displayable(Protocol):
+    def display(self, tab: str) -> str: ...
+
+def verbose_display(displayable_obj: _Displayable) -> None:
     if settings.logging_on_verbose:
         print(displayable_obj.display('\t'))
 
-def error(obj, level, message):
+@overload
+def error(obj: _Obj, level: Literal['fatal', 'Fatal'], message: str) -> NoReturn: ...
+@overload
+def error(obj: _Obj, level: str, message: str) -> None: ...
+
+def error(obj: _Obj, level: str, message: str) -> None:
     if settings.error_action == 'exit':
         print('{} {}: {}'.format(obj.__name__, level, message))
         if settings.fatal_traceback:

ubireader/py.typed

@@ -17,10 +17,18 @@
 # 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, log, verbose_log
 from ubireader.ubi.block import sort
 from ubireader.ubi.defines import UBI_VID_STATIC
 
+if TYPE_CHECKING:
+    from typing import Self # In TYPE_CHECKING block because it's only available after python 3.11
+    from collections.abc import Iterator, Mapping
+    from ubireader.ubi import ubi as Ubi
+    from ubireader.ubi.block import description as Block
+
 class ubi_file(object):
     """UBI image file object
 
@@ -50,7 +58,7 @@ class ubi_file(object):
     extract blocks, etc.
     """
 
-    def __init__(self, path, block_size, start_offset=0, end_offset=None):
+    def __init__(self, path: str, block_size: int, start_offset: int = 0, end_offset: int | None = None) -> None:
         self.__name__ = 'UBI_File'
         self.is_valid = False
         try:
@@ -89,54 +97,54 @@ def __init__(self, path, block_size, start_offset=0, end_offset=None):
         self._last_read_addr = self._fhandle.tell()
         self.is_valid = True
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
         self.close()
 
-    def _set_start(self, i):
+    def _set_start(self, i: int) -> None:
         self._start_offset = i
-    def _get_start(self):
+    def _get_start(self) -> int:
         return self._start_offset
     start_offset = property(_get_start, _set_start)
 
 
-    def _get_end(self):
+    def _get_end(self) -> int:
         return self._end_offset
     end_offset = property(_get_end)
 
 
-    def _get_block_size(self):
+    def _get_block_size(self) -> int:
         return self._block_size
     block_size = property(_get_block_size)
 
-    def close(self):
+    def close(self) -> None:
         self._fhandle.close()
 
-    def seek(self, offset):
+    def seek(self, offset: int) -> None:
         self._fhandle.seek(offset)
 
 
-    def read(self, size):
+    def read(self, size: int) -> bytes:
         self._last_read_addr = self.tell()
         verbose_log(self, 'read loc: %s, size: %s' % (self._last_read_addr, size))
         return self._fhandle.read(size)
 
 
-    def tell(self):
+    def tell(self) -> int:
         return self._fhandle.tell()
 
 
-    def last_read_addr(self):
+    def last_read_addr(self) -> int:
         return self._last_read_addr
 
 
-    def reset(self):
+    def reset(self) -> None:
         self._fhandle.seek(self.start_offset)
 
 
-    def reader(self):
+    def reader(self) -> Iterator[bytes]:
         self.reset()
         while True:
             cur_loc = self._fhandle.tell()
@@ -154,7 +162,7 @@ def reader(self):
             yield buf
 
 
-    def read_block(self, block):
+    def read_block(self, block: Block) -> bytes:
         """Read complete PEB data from file.
         
         Argument:
@@ -164,7 +172,7 @@ def read_block(self, block):
         return self._fhandle.read(block.size)
 
 
-    def read_block_data(self, block):
+    def read_block_data(self, block: Block) -> bytes:
         """Read LEB data from file
         
         Argument:
@@ -180,7 +188,7 @@ def read_block_data(self, block):
 
 
 class leb_virtual_file():
-    def __init__(self, ubi, block_list):
+    def __init__(self, ubi: Ubi, block_list: Mapping[int, Block]) -> None:
         self.__name__ = 'leb_virtual_file'
         self.is_valid = False
         self._ubi = ubi
@@ -196,7 +204,7 @@ def __init__(self, ubi, block_list):
             self.is_valid = True
 
 
-    def read(self, size):
+    def read(self, size: int) -> bytes:
         buf = ''
         leb = int(self.tell() / self._ubi.leb_size)
         offset = self.tell() % self._ubi.leb_size
@@ -227,24 +235,24 @@ def read(self, size):
                 error(self, 'Fatal', 'read loc: %s, size: %s, LEB: %s, offset: %s, error: %s' % (self._last_read_addr, size, leb, offset, e))
 
 
-    def reset(self):
+    def reset(self) -> None:
         self.seek(0)
 
 
-    def seek(self, offset):
+    def seek(self, offset: int) -> None:
         self._seek = offset
 
 
-    def tell(self):
+    def tell(self) -> int:
         return self._seek
 
 
-    def last_read_addr(self):
+    def last_read_addr(self) -> int:
         """Start address of last physical file read"""
         return self._last_read_addr
 
 
-    def reader(self):
+    def reader(self) -> Iterator[bytes]:
         last_leb = 0
         for block in self._blocks:
             while 0 != (self._ubi.blocks[block].leb_num - last_leb):

ubireader/ubi_io.py

@@ -17,13 +17,14 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
 import re
 from ubireader.debug import error, log
 from ubireader.ubi.defines import UBI_EC_HDR_MAGIC, FILE_CHUNK_SZ
 from ubireader.ubifs.defines import UBIFS_NODE_MAGIC, UBIFS_SB_NODE_SZ, UBIFS_SB_NODE, UBIFS_COMMON_HDR_SZ
 from ubireader.ubifs import nodes
 
-def guess_start_offset(path, guess_offset=0):
+def guess_start_offset(path: str, guess_offset: int =0) -> int | None:
     file_offset = guess_offset
 
     f = open(path, 'rb')
@@ -60,7 +61,7 @@ def guess_start_offset(path, guess_offset=0):
     f.close()
 
 
-def guess_filetype(path, start_offset=0):
+def guess_filetype(path: str, start_offset: int = 0) -> bytes | None:
     log(guess_filetype, 'Looking for file type at %s' % start_offset)
 
     with open(path, 'rb') as f:
@@ -81,7 +82,7 @@ def guess_filetype(path, start_offset=0):
     return ftype
 
 
-def guess_leb_size(path):
+def guess_leb_size(path: str) -> int | None:
     """Get LEB size from superblock
 
     Arguments:
@@ -125,7 +126,7 @@ def guess_leb_size(path):
     return block_size
 
 
-def guess_peb_size(path):
+def guess_peb_size(path: str) -> int | None:
     """Determine the most likely block size
 
     Arguments:
@@ -138,7 +139,7 @@ def guess_peb_size(path):
         common length between them.
     """
     file_offset = 0
-    offsets = []
+    offsets: list[int] = []
     f = open(path, 'rb')
     f.seek(0,2)
     file_size = f.tell()+1
@@ -160,7 +161,7 @@ def guess_peb_size(path):
         file_offset += FILE_CHUNK_SZ
     f.close()
 
-    occurrences = {}
+    occurrences: dict[int, int] = {}
     for i in range(0, len(offsets)):
         try:
             diff = offsets[i] - offsets[i-1]