Merge pull request #16 from quarkslab/refactoring_types

Make static typing PEP-585 compliant. NFC.

Author: RobinDavid

src/binexport/basic_block.py

@@ -1,15 +1,16 @@
+from __future__ import annotations
 import weakref
 from functools import cached_property
 from typing import TYPE_CHECKING
 
 from binexport.utils import instruction_index_range, get_instruction_address
 from binexport.instruction import InstructionBinExport
-from binexport.types import Addr
 
 if TYPE_CHECKING:
-    from .program import ProgramBinExport
-    from .function import FunctionBinExport
-    from .binexport2_pb2 import BinExport2
+    from binexport.program import ProgramBinExport
+    from binexport.function import FunctionBinExport
+    from binexport.binexport2_pb2 import BinExport2
+    from binexport.types import Addr
 
 
 class BasicBlockBinExport:
@@ -19,9 +20,9 @@ class BasicBlockBinExport:
 
     def __init__(
         self,
-        program: weakref.ref["ProgramBinExport"],
-        function: weakref.ref["FunctionBinExport"],
-        pb_bb: "BinExport2.BasicBlock",
+        program: weakref.ref[ProgramBinExport],
+        function: weakref.ref[FunctionBinExport],
+        pb_bb: BinExport2.BasicBlock,
     ):
         """
         :param program: Weak reference to the program
@@ -65,7 +66,7 @@ def __repr__(self) -> str:
         return "<%s:0x%x>" % (type(self).__name__, self.addr)
 
     @property
-    def program(self) -> "ProgramBinExport":
+    def program(self) -> ProgramBinExport:
         """
         Wrapper on weak reference on ProgramBinExport
 
@@ -74,7 +75,7 @@ def program(self) -> "ProgramBinExport":
         return self._program()
 
     @property
-    def function(self) -> "FunctionBinExport":
+    def function(self) -> FunctionBinExport:
         """
         Wrapper on weak reference on FunctionBinExport
 

src/binexport/expression.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import logging
 from functools import cached_property
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from binexport.binexport2_pb2 import BinExport2
 from binexport.types import ExpressionType
@@ -54,9 +54,9 @@ class ExpressionBinExport:
 
     def __init__(
         self,
-        program: "ProgramBinExport",
-        function: "FunctionBinExport",
-        instruction: "InstructionBinExport",
+        program: ProgramBinExport,
+        function: FunctionBinExport,
+        instruction: InstructionBinExport,
         exp_idx: int,
         parent: ExpressionBinExport | None = None,
     ):
@@ -70,7 +70,7 @@ def __init__(
         """
 
         self._idx = exp_idx
-        self.parent: Optional[ExpressionBinExport] = parent  #: parent expression if nested
+        self.parent: ExpressionBinExport | None = parent  #: parent expression if nested
         self.is_addr: bool = False  #: whether the value is referring to an address
         self.is_data: bool = False  #: whether the value is a reference to data
 
@@ -110,9 +110,9 @@ def depth(self) -> int:
 
     def _parse_protobuf(
         self,
-        program: "ProgramBinExport",
-        function: "FunctionBinExport",
-        instruction: "InstructionBinExport",
+        program: ProgramBinExport,
+        function: FunctionBinExport,
+        instruction: InstructionBinExport,
     ) -> None:
         """
         Low-level expression parser. It populates self._type and self._value

src/binexport/function.py

@@ -1,17 +1,19 @@
+from __future__ import annotations
 import logging
 import weakref
 import networkx
 from functools import cached_property
-from typing import TYPE_CHECKING, Dict, Set
+from typing import TYPE_CHECKING
 
 from binexport.utils import get_basic_block_addr
 from binexport.basic_block import BasicBlockBinExport
-from binexport.types import FunctionType, Addr
-from collections import abc
+from binexport.types import FunctionType
 
 if TYPE_CHECKING:
-    from .program import ProgramBinExport
-    from .binexport2_pb2 import BinExport2
+    from collections import abc
+    from binexport.program import ProgramBinExport
+    from binexport.binexport2_pb2 import BinExport2
+    from binexport.types import Addr
 
 
 class FunctionBinExport:
@@ -22,9 +24,9 @@ class FunctionBinExport:
 
     def __init__(
         self,
-        program: weakref.ref["ProgramBinExport"],
+        program: weakref.ref[ProgramBinExport],
         *,
-        pb_fun: "BinExport2.FlowGraph | None" = None,
+        pb_fun: BinExport2.FlowGraph | None = None,
         is_import: bool = False,
         addr: Addr | None = None,
     ):
@@ -40,8 +42,8 @@ def __init__(
         super(FunctionBinExport, self).__init__()
 
         self.addr: Addr | None = addr  #: address, None if imported function
-        self.parents: Set["FunctionBinExport"] = set()  #: set of function call this one
-        self.children: Set["FunctionBinExport"] = set()  #: set of functions called by this one
+        self.parents: set[FunctionBinExport] = set()  #: set of function call this one
+        self.children: set[FunctionBinExport] = set()  #: set of functions called by this one
 
         # Private attributes
         self._graph = None  # CFG. Loaded inside self.blocks
@@ -95,7 +97,7 @@ def unload(self) -> None:
         if self._enable_unloading:
             self._basic_blocks = None
 
-    def items(self) -> abc.ItemsView[Addr, "BasicBlockBinExport"]:
+    def items(self) -> abc.ItemsView[Addr, BasicBlockBinExport]:
         """
         Each function is associated to a dictionary with key-value
         Addr->BasicBlockBinExport. This returns items of the dictionary.
@@ -109,14 +111,14 @@ def keys(self) -> abc.KeysView[Addr]:
         """
         return self.blocks.keys()
 
-    def values(self) -> abc.ValuesView["BasicBlockBinExport"]:
+    def values(self) -> abc.ValuesView[BasicBlockBinExport]:
         """
         Each function is associated to a dictionary with key-value : Addr, BasicBlockBinExport. This returns items
         of the dictionary.
         """
         return self.blocks.values()
 
-    def __getitem__(self, item: Addr) -> "BasicBlockBinExport":
+    def __getitem__(self, item: Addr) -> BasicBlockBinExport:
         """
         Get a basic block object from its address.
 
@@ -135,14 +137,14 @@ def __contains__(self, item: Addr) -> bool:
         return item in self.blocks
 
     @property
-    def program(self) -> "ProgramBinExport":
+    def program(self) -> ProgramBinExport:
         """
         :py:class:`ProgramBinExport` in which this function belongs to.
         """
         return self._program()
 
     @property
-    def blocks(self) -> Dict[Addr, BasicBlockBinExport]:
+    def blocks(self) -> dict[Addr, BasicBlockBinExport]:
         """
         Returns a dict which is used to reference all basic blocks by their address.
         Calling this function will also load the CFG.

src/binexport/instruction.py

@@ -1,14 +1,15 @@
+from __future__ import annotations
 import weakref
 from functools import cached_property
-from typing import TYPE_CHECKING, List, Set
+from typing import TYPE_CHECKING
 
 from binexport.operand import OperandBinExport
-from binexport.types import Addr
 
 if TYPE_CHECKING:
     from .program import ProgramBinExport
     from .function import FunctionBinExport
     from .binexport2_pb2 import BinExport2
+    from binexport.types import Addr
 
 
 class InstructionBinExport:
@@ -18,8 +19,8 @@ class InstructionBinExport:
 
     def __init__(
         self,
-        program: weakref.ref["ProgramBinExport"],
-        function: weakref.ref["FunctionBinExport"],
+        program: weakref.ref[ProgramBinExport],
+        function: weakref.ref[FunctionBinExport],
         addr: Addr,
         i_idx: int,
     ):
@@ -33,7 +34,7 @@ def __init__(
         self._program = program
         self._function = function
         self._idx = i_idx
-        self.data_refs: Set[Addr] = self.program.data_refs[self._idx]  #: Data references address
+        self.data_refs: set[Addr] = self.program.data_refs[self._idx]  #: Data references address
         self.bytes = self.pb_instr.raw_bytes  #: bytes of the instruction (opcodes)
 
     def __hash__(self) -> int:
@@ -46,14 +47,14 @@ def __repr__(self) -> str:
         return f"<{type(self).__name__} {self.addr:#08x}: {self.mnemonic} {', '.join(str(x) for x in self.operands)}>"
 
     @property
-    def program(self) -> "ProgramBinExport":
+    def program(self) -> ProgramBinExport:
         """
         Program associated with this instruction.
         """
         return self._program()
 
     @property
-    def pb_instr(self) -> "BinExport2.Instruction":
+    def pb_instr(self) -> BinExport2.Instruction:
         """
         Protobuf instruction object.
         """
@@ -67,7 +68,7 @@ def mnemonic(self) -> str:
         return self.program.proto.mnemonic[self.pb_instr.mnemonic_index].name
 
     @cached_property
-    def operands(self) -> List[OperandBinExport]:
+    def operands(self) -> list[OperandBinExport]:
         """
         Returns a list of the operands instanciated dynamically on-demand.
         The list is cached by default, to erase the cache delete the attribute.

src/binexport/operand.py

@@ -1,11 +1,12 @@
-import weakref
+from __future__ import annotations
 from functools import cached_property
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 from binexport.expression import ExpressionBinExport
 from binexport.types import ExpressionType
 
 if TYPE_CHECKING:
+    import weakref
     from .program import ProgramBinExport
     from .function import FunctionBinExport
     from .instruction import InstructionBinExport
@@ -20,9 +21,9 @@ class OperandBinExport:
 
     def __init__(
         self,
-        program: weakref.ref["ProgramBinExport"],
-        function: weakref.ref["FunctionBinExport"],
-        instruction: weakref.ref["InstructionBinExport"],
+        program: weakref.ref[ProgramBinExport],
+        function: weakref.ref[FunctionBinExport],
+        instruction: weakref.ref[InstructionBinExport],
         op_idx: int,
     ):
         """
@@ -86,36 +87,36 @@ def __repr__(self) -> str:
         return f"<{type(self).__name__} {str(self)}>"
 
     @property
-    def program(self) -> "ProgramBinExport":
+    def program(self) -> ProgramBinExport:
         """
         Program object associated to this operand.
         """
         return self._program()
 
     @property
-    def function(self) -> "FunctionBinExport":
+    def function(self) -> FunctionBinExport:
         """
         Function object associated to this operand.
         """
 
         return self._function()
 
     @property
-    def instruction(self) -> "InstructionBinExport":
+    def instruction(self) -> InstructionBinExport:
         """
         Instruction object associated to this operand.
         """
         return self._instruction()
 
     @property
-    def pb_operand(self) -> "BinExport2.Operand":
+    def pb_operand(self) -> BinExport2.Operand:
         """
         Protobuf operand object in the protobuf structure.
         """
         return self.program.proto.operand[self._idx]
 
     @cached_property
-    def expressions(self) -> List[ExpressionBinExport]:
+    def expressions(self) -> list[ExpressionBinExport]:
         """
         Iterates over all the operand expression in a pre-order manner
         (binary operator first).

src/binexport/program.py

@@ -1,13 +1,17 @@
+from __future__ import annotations
 import pathlib
 import logging
 import networkx
 import weakref
 from collections import defaultdict
-from typing import Dict, Set
+from typing import TYPE_CHECKING
 
 from binexport.binexport2_pb2 import BinExport2
 from binexport.function import FunctionBinExport
-from binexport.types import FunctionType, Addr
+from binexport.types import FunctionType
+
+if TYPE_CHECKING:
+    from binexport.types import Addr
 
 
 class ProgramBinExport(dict):
@@ -27,12 +31,12 @@ def __init__(self, file: pathlib.Path | str):
         with open(file, "rb") as f:
             self._pb.ParseFromString(f.read())
         self.mask = 0xFFFFFFFF if self.architecture.endswith("32") else 0xFFFFFFFFFFFFFFFF
-        self.fun_names: Dict[str, "FunctionBinExport"] = {}  #: dictionary function name -> name
+        self.fun_names: dict[str, FunctionBinExport] = {}  #: dictionary function name -> name
         self.callgraph: networkx.DiGraph = networkx.DiGraph()  #: program callgraph (as Digraph)
 
         # Make the data refs map {instruction index -> address referred}
         # dictionary of instruction index to set of refs
-        self.data_refs: Dict[int, Set[Addr]] = defaultdict(set)
+        self.data_refs: dict[int, set[Addr]] = defaultdict(set)
         for entry in self.proto.data_reference:
             self.data_refs[entry.instruction_index].add(entry.address)
 
@@ -109,7 +113,7 @@ def from_binary_file(
         output_file: str | pathlib.Path = "",
         open_export: bool = True,
         override: bool = False,
-    ) -> "ProgramBinExport | bool":
+    ) -> ProgramBinExport | bool:
         """
         Generate the .BinExport file for the given program and return an instance
         of ProgramBinExport.

src/binexport/types.py

@@ -1,7 +1,7 @@
-from __future__ import annotations, absolute_import
+from __future__ import annotations
 import enum
-from typing import TypeAlias
 import enum_tools.documentation
+from typing import TypeAlias
 
 from binexport.binexport2_pb2 import BinExport2