From 7372aa91c6cb579d613cad163e26159bbb6d86bf Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Mon, 14 Jun 2021 10:56:44 -0600
Subject: [PATCH] engine: better type doc

---
 capa/engine.py          | 25 ++++++++++++++++++++-----
 capa/features/common.py |  9 +++------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/capa/engine.py b/capa/engine.py
index 0da6301c..e8b46b02 100644
--- a/capa/engine.py
+++ b/capa/engine.py
@@ -10,14 +10,17 @@ import copy
 import collections
 from typing import TYPE_CHECKING, Set, Dict, List, Tuple, Union, Mapping
 
-if TYPE_CHECKING:
-    from capa.rules import Rule
-
+import capa.rules
 import capa.features.common
 from capa.features.common import Feature
 
 # a collection of features and the locations at which they are found.
-# used throughout matching as the context in which features are searched.
+#
+# used throughout matching as the context in which features are searched:
+# to check if a feature exists, do: `Number(0x10) in features`.
+# to collect the locations of a feature, do: `features[Number(0x10)]`
+#
+# aliased here so that the type can be documented and xref'd.
 FeatureSet = Dict[Feature, Set[int]]
 
 
@@ -209,10 +212,22 @@ class Subscope(Statement):
 
 
 # mapping from rule name to list of: (location of match, result object)
+#
+# used throughout matching and rendering to collection the results
+#  of statement evaluation and their locations.
+#
+# to check if a rule matched, do: `"TCP client" in matches`.
+# to find where a rule matched, do: `map(first, matches["TCP client"])`
+# to see how a rule matched, do:
+#
+#     for address, match_details in matches["TCP client"]:
+#         inspect(match_details)
+#
+# aliased here so that the type can be documented and xref'd.
 MatchResults = Mapping[str, List[Tuple[int, Result]]]
 
 
-def match(rules: List["Rule"], features: FeatureSet, va: int) -> Tuple[FeatureSet, MatchResults]:
+def match(rules: List["capa.rules.Rule"], features: FeatureSet, va: int) -> Tuple[FeatureSet, MatchResults]:
     """
     Args:
       rules (List[capa.rules.Rule]): these must already be ordered topologically by dependency.
diff --git a/capa/features/common.py b/capa/features/common.py
index 806ac4ab..7e218b78 100644
--- a/capa/features/common.py
+++ b/capa/features/common.py
@@ -10,13 +10,10 @@ import re
 import codecs
 import logging
 import collections
-from typing import TYPE_CHECKING, Set, Dict, Union
+from typing import Set, Dict, Union
 
 import capa.engine
-import capa.features.common
-
-if TYPE_CHECKING:
-    from capa.engine import Result
+import capa.features
 
 logger = logging.getLogger(__name__)
 MAX_BYTES_FEATURE_SIZE = 0x100
@@ -104,7 +101,7 @@ class Feature:
     def __repr__(self):
         return str(self)
 
-    def evaluate(self, ctx: Dict["Feature", Set[int]]) -> "Result":
+    def evaluate(self, ctx: Dict["Feature", Set[int]]) -> "capa.engine.Result":
         return capa.engine.Result(self in ctx, self, [], locations=ctx.get(self, []))
 
     def freeze_serialize(self):