Add type hints

elftools/__init__.py

@@ -4,4 +4,4 @@
 # Eli Bendersky ([email protected])
 # This code is in the public domain
 #-------------------------------------------------------------------------------
-__version__ = '0.32'
+__version__: str = '0.32'

elftools/common/construct_utils.py

@@ -6,11 +6,20 @@
 # Eli Bendersky ([email protected])
 # This code is in the public domain
 #-------------------------------------------------------------------------------
+from __future__ import annotations
+
 from struct import Struct
+from typing import IO, TYPE_CHECKING, Any, NoReturn
+
 from ..construct import (
     Subconstruct, ConstructError, ArrayError, SizeofError, Construct, StaticField, FieldError
     )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from ..construct import Container
+
 
 class RepeatUntilExcluding(Subconstruct):
     """ A version of construct's RepeatUntil that doesn't include the last
@@ -21,12 +30,12 @@ class RepeatUntilExcluding(Subconstruct):
         P.S. removed some code duplication
     """
     __slots__ = ["predicate"]
-    def __init__(self, predicate, subcon):
+    def __init__(self, predicate: Callable[[Any, Container], bool], subcon: Construct) -> None:
         Subconstruct.__init__(self, subcon)
         self.predicate = predicate
         self._clear_flag(self.FLAG_COPY_CONTEXT)
         self._set_flag(self.FLAG_DYNAMIC)
-    def _parse(self, stream, context):
+    def _parse(self, stream: IO[bytes], context: Container) -> list[Any]:
         obj = []
         try:
             context_for_subcon = context
@@ -41,15 +50,17 @@ def _parse(self, stream, context):
         except ConstructError as ex:
             raise ArrayError("missing terminator", ex)
         return obj
-    def _build(self, obj, stream, context):
+    def _build(self, obj: Iterable[Any], stream: IO[bytes], context: Container) -> NoReturn:
         raise NotImplementedError('no building')
-    def _sizeof(self, context):
+    def _sizeof(self, context: Container) -> int:
         raise SizeofError("can't calculate size")
 
 class ULEB128(Construct):
     """A construct based parser for ULEB128 encoding.
     """
-    def _parse(self, stream, context):
+    if TYPE_CHECKING:
+        name: str  # instead of `str|None` from Construct to save us from `is None` checks everywhere
+    def _parse(self, stream: IO[bytes], context: Container) -> int:
         value = 0
         shift = 0
         while True:
@@ -65,7 +76,9 @@ def _parse(self, stream, context):
 class SLEB128(Construct):
     """A construct based parser for SLEB128 encoding.
     """
-    def _parse(self, stream, context):
+    if TYPE_CHECKING:
+        name: str  # instead of `str|None` from Construct to save us from `is None` checks everywhere
+    def _parse(self, stream: IO[bytes], context: Container) -> int:
         value = 0
         shift = 0
         while True:
@@ -88,40 +101,42 @@ class StreamOffset(Construct):
     Example:
     StreamOffset("item_offset")
     """
-    __slots__ = []
-    def __init__(self, name):
+    __slots__: list[str] = []
+    if TYPE_CHECKING:
+        name: str  # instead of `str|None` from Construct to save us from `is None` checks everywhere
+    def __init__(self, name: str) -> None:
         Construct.__init__(self, name)
         self._set_flag(self.FLAG_DYNAMIC)
-    def _parse(self, stream, context):
+    def _parse(self, stream: IO[bytes], context: Container) -> int:
         return stream.tell()
-    def _build(self, obj, stream, context):
+    def _build(self, obj: None, stream: IO[bytes], context: Container) -> None:
         context[self.name] = stream.tell()
-    def _sizeof(self, context):
+    def _sizeof(self, context: Container) -> int:
         return 0
 
 _UBInt24_packer = Struct(">BH")
 _ULInt24_packer = Struct("<HB")
 
 class UBInt24(StaticField):
     """unsigned, big endian 24-bit integer"""
-    def __init__(self, name):
+    def __init__(self, name: str) -> None:
         StaticField.__init__(self, name, 3)
 
-    def _parse(self, stream, context):
+    def _parse(self, stream: IO[bytes], context: Container) -> int:
         (h, l) = _UBInt24_packer.unpack(StaticField._parse(self, stream, context))
         return l | (h << 16)
 
-    def _build(self, obj, stream, context):
+    def _build(self, obj: int, stream: IO[bytes], context: Container) -> None:
         StaticField._build(self, _UBInt24_packer.pack(obj >> 16, obj & 0xFFFF), stream, context)
 
 class ULInt24(StaticField):
     """unsigned, little endian 24-bit integer"""
-    def __init__(self, name):
+    def __init__(self, name: str) -> None:
         StaticField.__init__(self, name, 3)
 
-    def _parse(self, stream, context):
+    def _parse(self, stream: IO[bytes], context: Container) -> int:
         (l, h) = _ULInt24_packer.unpack(StaticField._parse(self, stream, context))
         return l | (h << 16)
 
-    def _build(self, obj, stream, context):
+    def _build(self, obj: int, stream: IO[bytes], context: Container) -> None:
         StaticField._build(self, _ULInt24_packer.pack(obj & 0xFFFF, obj >> 16), stream, context)

elftools/common/utils.py

@@ -6,24 +6,43 @@
 # Eli Bendersky ([email protected])
 # This code is in the public domain
 #-------------------------------------------------------------------------------
+from __future__ import annotations
+
 from contextlib import contextmanager
+from typing import IO, TYPE_CHECKING, Any, TypeVar, overload
+
 from .exceptions import ELFParseError, ELFError, DWARFError
 from ..construct import ConstructError, ULInt8
 import os
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Mapping
+
+    from ..construct import Construct, FormatField
+    from ..dwarf.dwarfinfo import DebugSectionDescriptor
+    from .construct_utils import SLEB128, ULEB128, UBInt24, ULInt24
+
+    _T = TypeVar("_T")
+    _K = TypeVar("_K")
+    _V = TypeVar("_V")
+
 
-def merge_dicts(*dicts):
+def merge_dicts(*dicts: Mapping[_K, _V]) -> dict[_K, _V]:
     "Given any number of dicts, merges them into a new one."""
-    result = {}
+    result: dict[_K, _V] = {}
     for d in dicts:
         result.update(d)
     return result
 
-def bytes2str(b):
+def bytes2str(b: bytes) -> str:
     """Decode a bytes object into a string."""
     return b.decode('latin-1')
 
-def struct_parse(struct, stream, stream_pos=None):
+@overload
+def struct_parse(struct: FormatField[_T] | ULEB128 | SLEB128 | UBInt24 | ULInt24, stream: IO[bytes], stream_pos: int | None = ...) -> _T: ...
+@overload
+def struct_parse(struct: Construct, stream: IO[bytes], stream_pos: int | None = ...) -> Any: ...
+def struct_parse(struct: Construct, stream: IO[bytes], stream_pos: int | None = None) -> Any:
     """ Convenience function for using the given struct to parse a stream.
         If stream_pos is provided, the stream is seeked to this position before
         the parsing is done. Otherwise, the current position of the stream is
@@ -38,7 +57,7 @@ def struct_parse(struct, stream, stream_pos=None):
         raise ELFParseError(str(e))
 
 
-def parse_cstring_from_stream(stream, stream_pos=None):
+def parse_cstring_from_stream(stream: IO[bytes], stream_pos: int | None = None) -> bytes | None:
     """ Parse a C-string from the given stream. The string is returned without
         the terminating \x00 byte. If the terminating byte wasn't found, None
         is returned (the stream is exhausted).
@@ -67,20 +86,20 @@ def parse_cstring_from_stream(stream, stream_pos=None):
     return b''.join(chunks) if found else None
 
 
-def elf_assert(cond, msg=''):
+def elf_assert(cond: object, msg: str = '') -> None:
     """ Assert that cond is True, otherwise raise ELFError(msg)
     """
     _assert_with_exception(cond, msg, ELFError)
 
 
-def dwarf_assert(cond, msg=''):
+def dwarf_assert(cond: object, msg: str = '') -> None:
     """ Assert that cond is True, otherwise raise DWARFError(msg)
     """
     _assert_with_exception(cond, msg, DWARFError)
 
 
 @contextmanager
-def preserve_stream_pos(stream):
+def preserve_stream_pos(stream: IO[bytes]) -> Iterator[None]:
     """ Usage:
         # stream has some position FOO (return value of stream.tell())
         with preserve_stream_pos(stream):
@@ -92,18 +111,18 @@ def preserve_stream_pos(stream):
     stream.seek(saved_pos)
 
 
-def roundup(num, bits):
+def roundup(num: int, bits: int) -> int:
     """ Round up a number to nearest multiple of 2^bits. The result is a number
         where the least significant bits passed in bits are 0.
     """
     return (num - 1 | (1 << bits) - 1) + 1
 
-def read_blob(stream, length):
+def read_blob(stream: IO[bytes], length: int) -> list[int]:
     """Read length bytes from stream, return a list of ints
     """
     return [struct_parse(ULInt8(''), stream) for i in range(length)]
 
-def save_dwarf_section(section, filename):
+def save_dwarf_section(section: DebugSectionDescriptor, filename: str) -> None:
     """Debug helper: dump section contents into a file
     Section is expected to be one of the debug_xxx_sec elements of DWARFInfo
     """
@@ -116,21 +135,21 @@ def save_dwarf_section(section, filename):
         file.write(data)
     stream.seek(pos, os.SEEK_SET)
 
-def iterbytes(b):
+def iterbytes(b: bytes) -> Iterator[bytes]:
     """Return an iterator over the elements of a bytes object.
 
     For example, for b'abc' yields b'a', b'b' and then b'c'.
     """
     for i in range(len(b)):
         yield b[i:i+1]
 
-def bytes2hex(b, sep=''):
+def bytes2hex(b: bytes, sep: str = '') -> str:
     if not sep:
         return b.hex()
     return sep.join(map('{:02x}'.format, b))
 
 #------------------------- PRIVATE -------------------------
 
-def _assert_with_exception(cond, msg, exception_type):
+def _assert_with_exception(cond: object, msg: str, exception_type: type[BaseException]) -> None:
     if not cond:
         raise exception_type(msg)

elftools/construct/lib/container.py

@@ -65,7 +65,7 @@ def __getitem__(self, name: Literal[
         "length",
         "n_descsz", "n_offset", "n_namesz",
         "sh_addralign", "sh_flags", "sh_size",
-        "bloom_size", "nbuckets",
+        "bloom_size", "nbuckets", "nchains",
     ]) -> int: ...
     @overload
     def __getitem__(self, name: Literal[
@@ -75,6 +75,10 @@ def __getitem__(self, name: Literal[
         "tag", "vendor_name",
     ]) -> str: ...
     @overload
+    def __getitem__(self, name: Literal[
+        "buckets", "chains",
+    ]) -> list[int]: ...
+    @overload
     def __getitem__(self, name: str) -> Any: ...
     def __getitem__(self, name: str) -> Any:
         return self.__dict__[name]

elftools/construct/lib/hex.py

@@ -7,8 +7,7 @@
 
 
 # Map an integer in the inclusive range 0-255 to its string byte representation
-_printable = dict((i, ".") for i in range(256))
-_printable.update((i, chr(i)) for i in range(32, 128))
+_printable = {i: chr(i) if 32 <= i < 128 else "." for i in range(256)}
 
 
 def hexdump(data: bytes, linesize: int) -> list[str]:

elftools/dwarf/abbrevtable.py

@@ -6,14 +6,24 @@
 # Eli Bendersky ([email protected])
 # This code is in the public domain
 #-------------------------------------------------------------------------------
+from __future__ import annotations
+
+from typing import IO, TYPE_CHECKING, Any
+
 from ..common.utils import struct_parse
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from ..construct.lib.container import Container
+    from .structs import DWARFStructs
+
 
 class AbbrevTable:
     """ Represents a DWARF abbreviation table.
     """
     __slots__ = ('structs', 'stream', 'offset', '_abbrev_map')
-    def __init__(self, structs, stream, offset):
+    def __init__(self, structs: DWARFStructs, stream: IO[bytes], offset: int) -> None:
         """ Create new abbreviation table. Parses the actual table from the
             stream and stores it internally.
 
@@ -30,19 +40,19 @@ def __init__(self, structs, stream, offset):
 
         self._abbrev_map = self._parse_abbrev_table()
 
-    def get_abbrev(self, code):
+    def get_abbrev(self, code: int) -> AbbrevDecl:
         """ Get the AbbrevDecl for a given code. Raise KeyError if no
             declaration for this code exists.
         """
         return self._abbrev_map[code]
 
-    def _parse_abbrev_table(self):
+    def _parse_abbrev_table(self) -> dict[int, AbbrevDecl]:
         """ Parse the abbrev table from the stream
         """
-        map = {}
+        map: dict[int, AbbrevDecl] = {}
         self.stream.seek(self.offset)
         while True:
-            decl_code = struct_parse(
+            decl_code: int = struct_parse(
                 struct=self.structs.the_Dwarf_uleb128,
                 stream=self.stream)
             if decl_code == 0:
@@ -61,20 +71,20 @@ class AbbrevDecl:
         The abbreviation declaration represents an "entry" that points to it.
     """
     __slots__ = ('code', 'decl', '_has_children')
-    def __init__(self, code, decl):
+    def __init__(self, code: int, decl: Container) -> None:
         self.code = code
         self.decl = decl
         self._has_children = decl['children_flag'] == 'DW_CHILDREN_yes'
 
-    def has_children(self):
+    def has_children(self) -> bool:
         return self._has_children
 
-    def iter_attr_specs(self):
+    def iter_attr_specs(self) -> Iterator[tuple[str, str]]:
         """ Iterate over the attribute specifications for the entry. Yield
             (name, form) pairs.
         """
         for attr_spec in self['attr_spec']:
             yield attr_spec.name, attr_spec.form
 
-    def __getitem__(self, entry):
+    def __getitem__(self, entry: str) -> Any:
         return self.decl[entry]