fix: type annotations for disable_progress and module attribute access

capa/capabilities/common.py

@@ -64,7 +64,9 @@ class Capabilities:
     library_functions: Optional[tuple[LibraryFunction, ...]] = None
 
 
-def find_capabilities(ruleset: RuleSet, extractor: FeatureExtractor, disable_progress=None, **kwargs) -> Capabilities:
+def find_capabilities(
+    ruleset: RuleSet, extractor: FeatureExtractor, disable_progress: Optional[bool] = None, **kwargs
+) -> Capabilities:
     from capa.capabilities.static import find_static_capabilities
     from capa.capabilities.dynamic import find_dynamic_capabilities
 

capa/capabilities/dynamic.py

@@ -15,6 +15,7 @@
 import logging
 import itertools
 import collections
+from typing import Optional
 from dataclasses import dataclass
 
 import capa.perf
@@ -269,7 +270,7 @@ def find_process_capabilities(
 
 
 def find_dynamic_capabilities(
-    ruleset: RuleSet, extractor: DynamicFeatureExtractor, disable_progress: bool = False
+    ruleset: RuleSet, extractor: DynamicFeatureExtractor, disable_progress: Optional[bool] = None
 ) -> Capabilities:
     all_process_matches: MatchResults = collections.defaultdict(list)
     all_thread_matches: MatchResults = collections.defaultdict(list)
@@ -285,7 +286,7 @@ def find_dynamic_capabilities(
     n_processes: int = len(processes)
 
     with capa.helpers.CapaProgressBar(
-        console=capa.helpers.log_console, transient=True, disable=disable_progress
+        console=capa.helpers.log_console, transient=True, disable=bool(disable_progress)
     ) as pbar:
         task = pbar.add_task("matching", total=n_processes, unit="processes")
         for p in processes:

capa/capabilities/static.py

@@ -16,6 +16,7 @@ import time
 import logging
 import itertools
 import collections
+from typing import Optional
 from dataclasses import dataclass
 
 import capa.perf
@@ -149,7 +150,7 @@ def find_code_capabilities(ruleset: RuleSet, extractor: StaticFeatureExtractor,
 
 
 def find_static_capabilities(
-    ruleset: RuleSet, extractor: StaticFeatureExtractor, disable_progress=None
+    ruleset: RuleSet, extractor: StaticFeatureExtractor, disable_progress: Optional[bool] = None
 ) -> Capabilities:
     all_function_matches: MatchResults = collections.defaultdict(list)
     all_bb_matches: MatchResults = collections.defaultdict(list)
@@ -168,7 +169,7 @@ def find_static_capabilities(
     percentage: float = 0
 
     with capa.helpers.CapaProgressBar(
-        console=capa.helpers.log_console, transient=True, disable=disable_progress
+        console=capa.helpers.log_console, transient=True, disable=bool(disable_progress)
     ) as pbar:
         task = pbar.add_task(
             "matching", total=n_funcs, unit="functions", postfix=f"skipped {n_libs} library functions, {percentage}%"

capa/features/extractors/common.py

@@ -30,9 +30,11 @@ from capa.features.common import (
     OS_ANY,
     OS_AUTO,
     ARCH_ANY,
+    VALID_OS,
     FORMAT_PE,
     FORMAT_ELF,
     OS_WINDOWS,
+    VALID_ARCH,
     FORMAT_FREEZE,
     FORMAT_RESULT,
     Arch,
@@ -52,7 +54,7 @@ MATCH_RESULT = b'{"meta":'
 MATCH_JSON_OBJECT = b'{"'
 
 
-def extract_file_strings(buf: bytes, **kwargs) -> Iterator[tuple[String, Address]]:
+def extract_file_strings(buf: bytes) -> Iterator[tuple[String, Address]]:
     """
     extract ASCII and UTF-16 LE strings from file
     """
@@ -97,7 +99,7 @@ def extract_arch(buf) -> Iterator[tuple[Feature, Address]]:
         with contextlib.closing(io.BytesIO(buf)) as f:
             arch = capa.features.extractors.elf.detect_elf_arch(f)
 
-        if arch not in capa.features.common.VALID_ARCH:
+        if arch not in VALID_ARCH:
             logger.debug("unsupported arch: %s", arch)
             return
 
@@ -130,7 +132,7 @@ def extract_os(buf, os=OS_AUTO) -> Iterator[tuple[Feature, Address]]:
         with contextlib.closing(io.BytesIO(buf)) as f:
             os = capa.features.extractors.elf.detect_elf_os(f)
 
-        if os not in capa.features.common.VALID_OS:
+        if os not in VALID_OS:
             logger.debug("unsupported os: %s", os)
             return
 

tests/test_capabilities.py

@@ -226,7 +226,7 @@ def test_com_feature_matching(z395eb_extractor):
                     """)
         )
     ])
-    capabilities = capa.main.find_capabilities(rules, z395eb_extractor)
+    capabilities = capa.capabilities.common.find_capabilities(rules, z395eb_extractor)
     assert "initialize IWebBrowser2" in capabilities.matches