From 9b7e67443b790cc100c01985ad01a134b9dbde4a Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Mon, 14 Jun 2021 08:59:23 -0600
Subject: [PATCH] extractors: fix type hints

---
 capa/features/extractors/base_extractor.py  | 2 +-
 capa/features/extractors/smda/basicblock.py | 2 +-
 capa/features/extractors/smda/function.py   | 2 +-
 capa/features/extractors/smda/insn.py       | 2 +-
 capa/features/extractors/viv/basicblock.py  | 2 +-
 capa/features/extractors/viv/function.py    | 2 +-
 capa/features/extractors/viv/insn.py        | 2 +-
 7 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/capa/features/extractors/base_extractor.py b/capa/features/extractors/base_extractor.py
index 4808a656..67ba9451 100644
--- a/capa/features/extractors/base_extractor.py
+++ b/capa/features/extractors/base_extractor.py
@@ -171,7 +171,7 @@ class FeatureExtractor:
         raise NotImplemented
 
     @abc.abstractmethod
-    def get_instructions(self: FunctionHandle, f, bb: BBHandle) -> Iterator[InsnHandle]:
+    def get_instructions(self, f: FunctionHandle, bb: BBHandle) -> Iterator[InsnHandle]:
         """
         enumerate the instructions in the given basic block and provide opaque values that will
          subsequently be provided to `.extract_insn_features()`, etc.
diff --git a/capa/features/extractors/smda/basicblock.py b/capa/features/extractors/smda/basicblock.py
index b68d02a0..840712ce 100644
--- a/capa/features/extractors/smda/basicblock.py
+++ b/capa/features/extractors/smda/basicblock.py
@@ -116,7 +116,7 @@ def extract_features(f, bb):
       bb (smda.common.SmdaBasicBlock): the basic block to process.
 
     yields:
-      Feature, set[VA]: the features and their location found in this basic block.
+      Tuple[Feature, int]: the features and their location found in this basic block.
     """
     yield BasicBlock(), bb.offset
     for bb_handler in BASIC_BLOCK_HANDLERS:
diff --git a/capa/features/extractors/smda/function.py b/capa/features/extractors/smda/function.py
index 33c6c5d1..0308a29c 100644
--- a/capa/features/extractors/smda/function.py
+++ b/capa/features/extractors/smda/function.py
@@ -28,7 +28,7 @@ def extract_features(f):
       f (smda.common.SmdaFunction): the function from which to extract features
 
     yields:
-      Feature, set[VA]: the features and their location found in this function.
+      Tuple[Feature, int]: the features and their location found in this function.
     """
     for func_handler in FUNCTION_HANDLERS:
         for feature, va in func_handler(f):
diff --git a/capa/features/extractors/smda/insn.py b/capa/features/extractors/smda/insn.py
index eddc859d..da5ebbfa 100644
--- a/capa/features/extractors/smda/insn.py
+++ b/capa/features/extractors/smda/insn.py
@@ -370,7 +370,7 @@ def extract_features(f, bb, insn):
       insn (smda.common.SmdaInstruction): the instruction to process.
 
     yields:
-      Feature, set[VA]: the features and their location found in this insn.
+      Tuple[Feature, int]: the features and their location found in this insn.
     """
     for insn_handler in INSTRUCTION_HANDLERS:
         for feature, va in insn_handler(f, bb, insn):
diff --git a/capa/features/extractors/viv/basicblock.py b/capa/features/extractors/viv/basicblock.py
index ba8cf09c..39c2d2a4 100644
--- a/capa/features/extractors/viv/basicblock.py
+++ b/capa/features/extractors/viv/basicblock.py
@@ -152,7 +152,7 @@ def extract_features(f, bb):
       bb (viv_utils.BasicBlock): the basic block to process.
 
     yields:
-      Feature, set[VA]: the features and their location found in this basic block.
+      Tuple[Feature, int]: the features and their location found in this basic block.
     """
     yield BasicBlock(), bb.va
     for bb_handler in BASIC_BLOCK_HANDLERS:
diff --git a/capa/features/extractors/viv/function.py b/capa/features/extractors/viv/function.py
index 5426d569..60995868 100644
--- a/capa/features/extractors/viv/function.py
+++ b/capa/features/extractors/viv/function.py
@@ -61,7 +61,7 @@ def extract_features(f):
       f (viv_utils.Function): the function from which to extract features
 
     yields:
-      Feature, set[VA]: the features and their location found in this function.
+      Tuple[Feature, int]: the features and their location found in this function.
     """
     for func_handler in FUNCTION_HANDLERS:
         for feature, va in func_handler(f):
diff --git a/capa/features/extractors/viv/insn.py b/capa/features/extractors/viv/insn.py
index fb1b2d45..88515fb6 100644
--- a/capa/features/extractors/viv/insn.py
+++ b/capa/features/extractors/viv/insn.py
@@ -607,7 +607,7 @@ def extract_features(f, bb, insn):
       insn (vivisect...Instruction): the instruction to process.
 
     yields:
-      Feature, set[VA]: the features and their location found in this insn.
+      Tuple[Feature, int]: the features and their location found in this insn.
     """
     for insn_handler in INSTRUCTION_HANDLERS:
         for feature, va in insn_handler(f, bb, insn):