fix: lint driven fixes

Author: evilsocket

dyana/loaders/automodel/main.py

@@ -6,7 +6,7 @@
 import torch
 from transformers import AutoModel, AutoTokenizer
 
-from dyana import get_current_imports, get_gpu_usage, get_peak_rss
+from dyana import get_current_imports, get_gpu_usage, get_peak_rss  # type: ignore[attr-defined]
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Profile model files")

dyana/loaders/base/dyana.py

@@ -6,7 +6,7 @@
 
 
 @contextmanager
-def capture_output():
+def capture_output() -> t.Generator[tuple[StringIO, StringIO], None, None]:
     """
     Context manager to capture stdout and stderr
 
@@ -63,11 +63,11 @@ def get_gpu_usage() -> list[dict[str, t.Any]]:
     return usage
 
 
-def get_current_imports() -> dict[str, str]:
+def get_current_imports() -> dict[str, str | None]:
     """
     Get the currently imported modules.
     """
-    imports = {}
+    imports: dict[str, str | None] = {}
 
     # for each loaded module
     for module_name, module in sys.modules.items():

dyana/loaders/elf/main.py

@@ -4,7 +4,7 @@
 import subprocess
 import typing as t
 
-from dyana import get_peak_rss
+from dyana import get_peak_rss  # type: ignore[attr-defined]
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run an ELF file")

dyana/loaders/pickle/main.py

@@ -4,7 +4,7 @@
 import pickle
 import typing as t
 
-from dyana import get_current_imports, get_peak_rss
+from dyana import get_current_imports, get_peak_rss  # type: ignore[attr-defined]
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run a pickle file")

dyana/loaders/python/main.py

@@ -4,7 +4,7 @@
 import runpy
 import typing as t
 
-from dyana import capture_output, get_current_imports, get_peak_rss
+from dyana import capture_output, get_current_imports, get_peak_rss  # type: ignore[attr-defined]
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run an Python file")

dyana/view.py

@@ -267,7 +267,7 @@ def view_extra_unknown(key: str, value: t.Any) -> None:
 def count_package_prefixes(path_dict: dict[str, str], level: int = 2) -> dict[str, int]:
     from collections import defaultdict
 
-    prefix_counter = defaultdict(int)
+    prefix_counter: defaultdict[str, int] = defaultdict(int)
 
     for package_path in path_dict.keys():
         parts = package_path.split(".")