fix TypeAlias importing: import from typing_extensions to support Python 3.9 and lower

capa/features/extractors/base_extractor.py

@@ -8,9 +8,11 @@
 
 import abc
 import dataclasses
-from typing import Any, Dict, Tuple, Union, Iterator, TypeAlias
+from typing import Any, Dict, Tuple, Union, Iterator
 from dataclasses import dataclass
 
+from typing_extensions import TypeAlias
+
 import capa.features.address
 from capa.features.common import Feature
 from capa.features.address import Address, AbsoluteVirtualAddress

capa/main.py

@@ -231,14 +231,7 @@ def find_code_capabilities(
 def find_file_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, function_features: FeatureSet):
     file_features = collections.defaultdict(set)  # type: FeatureSet
 
-    if isinstance(extractor, StaticFeatureExtractor):
-        extractor_: StaticFeatureExtractor = cast(StaticFeatureExtractor, extractor)
-    elif isinstance(extractor, DynamicFeatureExtractor):
-        extractor_: DynamicFeatureExtractor = cast(DynamicFeatureExtractor, extractor)
-    else:
-        raise ValueError(f"unexpected extractor type: {extractor.__class__.__name__}")
-
-    for feature, va in itertools.chain(extractor_.extract_file_features(), extractor_.extract_global_features()):
+    for feature, va in itertools.chain(extractor.extract_file_features(), extractor.extract_global_features()):
         # not all file features may have virtual addresses.
         # if not, then at least ensure the feature shows up in the index.
         # the set of addresses will still be empty.
@@ -1251,7 +1244,8 @@ def main(argv=None):
         if format_ == FORMAT_FREEZE:
             # freeze format deserializes directly into an extractor
             with open(args.sample, "rb") as f:
-                extractor = frz.load(f.read())
+                extractor: FeatureExtractor = frz.load(f.read())
+                assert isinstance(extractor, StaticFeatureExtractor)
         else:
             # all other formats we must create an extractor,
             # such as viv, binary ninja, etc. workspaces
@@ -1270,7 +1264,7 @@ def main(argv=None):
             should_save_workspace = os.environ.get("CAPA_SAVE_WORKSPACE") not in ("0", "no", "NO", "n", None)
 
             try:
-                extractor: FeatureExtractor = get_extractor(
+                extractor = get_extractor(
                     args.sample,
                     format_,
                     args.os,

scripts/profile-time.py

@@ -46,7 +46,7 @@ import capa.helpers
 import capa.features
 import capa.features.common
 import capa.features.freeze
-from capa.features.extractors.base_extractor import FeatureExtractor
+from capa.features.extractors.base_extractor import StaticFeatureExtractor
 
 logger = logging.getLogger("capa.profile")
 
@@ -105,8 +105,9 @@ def main(argv=None):
     ):
         with open(args.sample, "rb") as f:
             extractor = capa.features.freeze.load(f.read())
+            assert isinstance(extractor, StaticFeatureExtractor)
     else:
-        extractor: FeatureExtractor = capa.main.get_extractor(
+        extractor = capa.main.get_extractor(
             args.sample, args.format, args.os, capa.main.BACKEND_VIV, sig_paths, should_save_workspace=False
         )