type: capa.render

capa/main.py

@@ -174,11 +174,10 @@ def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, disable_pro
 
     # collection of features that captures the rule matches within function and BB scopes.
     # mapping from feature (matched rule) to set of addresses at which it matched.
-    # schema: Dict[MatchedRule: Set[int]
     function_and_lower_features = {
         capa.features.common.MatchedRule(rule_name): set(map(lambda p: p[0], results))
         for rule_name, results in itertools.chain(all_function_matches.items(), all_bb_matches.items())
-    }
+    }  # type: FeatureSet
 
     all_file_matches, feature_count = find_file_capabilities(ruleset, extractor, function_and_lower_features)
     meta["feature_counts"]["file"] = feature_count

capa/render/default.py

@@ -7,14 +7,14 @@
 # See the License for the specific language governing permissions and limitations under the License.
 
 import collections
-from typing import Dict, List, Tuple
+from typing import Dict, List
 
 import tabulate
 
 import capa.render.utils as rutils
 import capa.render.result_document
 from capa.rules import RuleSet
-from capa.engine import Result
+from capa.engine import MatchResults
 from capa.render.utils import StringIO
 
 tabulate.PRESERVE_WHITESPACE = True
@@ -226,6 +226,6 @@ def render_default(doc):
     return ostream.getvalue()
 
 
-def render(meta, rules: RuleSet, capabilities: Dict[str, List[Tuple[int, Result]]]) -> str:
+def render(meta, rules: RuleSet, capabilities: MatchResults) -> str:
     doc = capa.render.result_document.convert_capabilities_to_result_document(meta, rules, capabilities)
     return render_default(doc)

capa/render/json.py

@@ -8,6 +8,8 @@
 import json
 
 import capa.render.result_document
+from capa.rules import RuleSet
+from capa.engine import MatchResults
 
 
 class CapaJsonObjectEncoder(json.JSONEncoder):
@@ -23,7 +25,7 @@ class CapaJsonObjectEncoder(json.JSONEncoder):
             return json.JSONEncoder.default(self, obj)
 
 
-def render(meta, rules, capabilities):
+def render(meta, rules: RuleSet, capabilities: MatchResults) -> str:
     return json.dumps(
         capa.render.result_document.convert_capabilities_to_result_document(meta, rules, capabilities),
         cls=CapaJsonObjectEncoder,

capa/render/result_document.py

@@ -12,7 +12,7 @@ import capa.engine
 import capa.render.utils
 import capa.features.common
 from capa.rules import RuleSet
-from capa.engine import Result
+from capa.engine import MatchResults
 
 
 def convert_statement_to_result_document(statement):
@@ -261,7 +261,7 @@ def parse_canonical_mbc(mbc: str):
     }
 
 
-def convert_capabilities_to_result_document(meta, rules: RuleSet, capabilities: Dict[str, List[Tuple[int, Result]]]):
+def convert_capabilities_to_result_document(meta, rules: RuleSet, capabilities: MatchResults):
     """
     convert the given rule set and capabilities result to a common, Python-native data structure.
     this format can be directly emitted to JSON, or passed to the other `capa.render.*.render()` routines

capa/render/verbose.py

@@ -27,6 +27,8 @@ import tabulate
 import capa.rules
 import capa.render.utils as rutils
 import capa.render.result_document
+from capa.rules import RuleSet
+from capa.engine import MatchResults
 
 
 def render_meta(ostream, doc):
@@ -123,6 +125,6 @@ def render_verbose(doc):
     return ostream.getvalue()
 
 
-def render(meta, rules, capabilities):
+def render(meta, rules: RuleSet, capabilities: MatchResults) -> str:
     doc = capa.render.result_document.convert_capabilities_to_result_document(meta, rules, capabilities)
     return render_verbose(doc)

capa/render/vverbose.py

@@ -13,6 +13,8 @@ import capa.render.utils as rutils
 import capa.render.verbose
 import capa.features.common
 import capa.render.result_document
+from capa.rules import RuleSet
+from capa.engine import MatchResults
 
 
 def render_locations(ostream, match):
@@ -264,6 +266,6 @@ def render_vverbose(doc):
     return ostream.getvalue()
 
 
-def render(meta, rules, capabilities):
+def render(meta, rules: RuleSet, capabilities: MatchResults) -> str:
     doc = capa.render.result_document.convert_capabilities_to_result_document(meta, rules, capabilities)
     return render_vverbose(doc)