diff --git a/capa/engine.py b/capa/engine.py index f496a146..2b4c4416 100644 --- a/capa/engine.py +++ b/capa/engine.py @@ -9,7 +9,13 @@ import copy import collections +from typing import Union, List, Dict, Set +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from capa.rules import Rule + import capa.features.common +from capa.features.common import Feature class Statement(object): @@ -78,7 +84,7 @@ class Result(object): we need this so that we can render the tree of expressions and their results. """ - def __init__(self, success, statement, children, locations=None): + def __init__(self, success: bool, statement: Union[Statement, Feature], children: List["Result"], locations=None): """ args: success (bool) @@ -199,7 +205,7 @@ class Subscope(Statement): raise ValueError("cannot evaluate a subscope directly!") -def match(rules, features, va): +def match(rules: List["Rule"], features: Dict[Feature, Set[int]], va: int): """ Args: rules (List[capa.rules.Rule]): these must already be ordered topologically by dependency. diff --git a/capa/rules.py b/capa/rules.py index fb2eb877..a8a5d2e7 100644 --- a/capa/rules.py +++ b/capa/rules.py @@ -33,7 +33,6 @@ import capa.features.file import capa.features.insn import capa.features.common import capa.features.basicblock -from capa.rules import Rule from capa.features.common import MAX_BYTES_FEATURE_SIZE, Feature logger = logging.getLogger(__name__) diff --git a/setup.py b/setup.py index b0e3a496..11f68e72 100644 --- a/setup.py +++ b/setup.py @@ -25,6 +25,7 @@ requirements = [ "vivisect==1.0.3", "smda==1.5.18", "pefile==2021.5.24", + "typing==3.7.4.3", ] # this sets __version__