FeatureExtractor alias: fix mypy typing issues by adding ininstance-based assert statements

This commit is contained in:
Yacine Elhamer
2023-06-26 22:46:27 +01:00
parent 63e4d3d5eb
commit b172f9a354
4 changed files with 15 additions and 13 deletions

View File

@@ -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 StaticFeatureExtractor
from capa.features.extractors.base_extractor import FeatureExtractor, StaticFeatureExtractor
logger = logging.getLogger("capa.profile")
@@ -104,13 +104,14 @@ def main(argv=None):
args.format == capa.features.common.FORMAT_AUTO and capa.features.freeze.is_freeze(taste)
):
with open(args.sample, "rb") as f:
extractor = capa.features.freeze.load(f.read())
extractor: FeatureExtractor = capa.features.freeze.load(f.read())
assert isinstance(extractor, StaticFeatureExtractor)
else:
extractor = capa.main.get_extractor(
args.sample, args.format, args.os, capa.main.BACKEND_VIV, sig_paths, should_save_workspace=False
)
assert isinstance(extractor, StaticFeatureExtractor)
with tqdm.tqdm(total=args.number * args.repeat) as pbar:
def do_iteration():