[mypy] add suggested type hints and fix some old ones

4 files changed, 26 insertions, 18 deletions

diff --git a/allocbench/allocator.py b/allocbench/allocator.py
index 8428580..8e4f041 100644
--- a/allocbench/allocator.py
+++ b/allocbench/allocator.py
@@ -25,8 +25,9 @@ from pathlib import Path
 import shutil
 from subprocess import CalledProcessError
 import sys
+from typing import List, Optional
 
-from allocbench.artifact import ArchiveArtifact, GitArtifact
+from allocbench.artifact import Artifact, ArchiveArtifact, GitArtifact
 from allocbench.globalvars import ALLOCBUILDDIR, ALLOCSRCDIR
 from allocbench.util import print_status, print_debug, print_error, print_info2, run_cmd
 
@@ -56,16 +57,16 @@ class Allocator:
         "prepare_cmds", "build_cmds"
     ]
 
-    binary_suffix = None
-    cmd_prefix = None
-    ld_preload = None
-    ld_library_path = None
+    binary_suffix: Optional[str] = None
+    cmd_prefix: Optional[str] = None
+    ld_preload: Optional[str] = None
+    ld_library_path: Optional[str] = None
     color = None
-    sources = None
-    version = None
-    patches = []
-    prepare_cmds = []
-    build_cmds = []
+    sources: Optional[Artifact] = None
+    version: Optional[str] = None
+    patches: List[str] = []
+    prepare_cmds: List[str] = []
+    build_cmds: List[str] = []
     analyze_alloc = False
 
     def __init__(self, name, **kwargs):
diff --git a/allocbench/benchmark.py b/allocbench/benchmark.py
index 1d75d67..923b84f 100644
--- a/allocbench/benchmark.py
+++ b/allocbench/benchmark.py
@@ -29,6 +29,7 @@ import os
 import subprocess
 from time import sleep
 import traceback
+from typing import Dict, Iterable, List, Optional
 
 import numpy as np
 
@@ -46,10 +47,10 @@ class Benchmark:
     perf_allowed = None
 
     cmd = "false"
-    args = {}
+    args: Dict[str, Iterable] = {}
     measure_cmd_csv = False
     measure_cmd = "perf stat -x, -d"
-    servers = []
+    servers: List[Dict[str, str]] = []
 
     @staticmethod
     def terminate_subprocess(proc, timeout=5):
@@ -736,7 +737,7 @@ class Benchmark:
                 self.results["stats"][alloc][perm] = stats
 
 
-def get_benchmark_object(benchmark_name: str) -> Benchmark:
+def get_benchmark_object(benchmark_name: str) -> Optional[Benchmark]:
     """Find the first Benchmark class in allocbench.benchmarks.{benchmark_name} and return an instance"""
     bench_module = importlib.import_module(
         f"allocbench.benchmarks.{benchmark_name}")
@@ -747,3 +748,5 @@ def get_benchmark_object(benchmark_name: str) -> Benchmark:
             continue
 
         return member()
+
+    return None
diff --git a/allocbench/globalvars.py b/allocbench/globalvars.py
index f2a98f0..1b3e5fb 100644
--- a/allocbench/globalvars.py
+++ b/allocbench/globalvars.py
@@ -31,10 +31,13 @@ RESDIR: Directory were the benchmark results are stored
 
 import inspect
 import os
+from typing import Dict
+
+import allocbench.allocator
 
 VERBOSITY = 0
 
-ALLOCATORS = {}
+ALLOCATORS: Dict[str, allocbench.allocator.Allocator] = {}
 
 # /.../allocbench/allocbench
 SRCDIR = os.path.dirname(
diff --git a/allocbench/plots.py b/allocbench/plots.py
index a55f7e2..f172355 100644
--- a/allocbench/plots.py
+++ b/allocbench/plots.py
@@ -22,13 +22,13 @@ import itertools
 import operator
 import os
 import re
-import scipy.stats
 import traceback
 from typing import Dict, List, Tuple, NamedTuple
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+import scipy.stats
 
 import allocbench.facter as facter
 from allocbench.globalvars import SUMMARY_FILE_EXT, LATEX_CUSTOM_PREAMBLE
@@ -148,7 +148,8 @@ def get_y_data(bench,
                stat="mean",
                scale=None) -> List[float]:
     """Helper to get the y data of an allocator for given permutations"""
-    y_data = []
+
+    y_data: List[float] = []
 
     if isinstance(perms, bench.Perm):
         perms = [perms]
@@ -552,9 +553,9 @@ def export_stats_to_csv(bench, datapoint, path=None):
 def get_ordered_results_for_perm(bench,
                                  datapoint: str,
                                  perm: NamedTuple,
-                                 order='>') -> List[Tuple[float, List]]:
+                                 order='>') -> List[Tuple[float, List[str]]]:
     """Return a ordered list of the allocator and their results for a specific perm"""
-    data = {}
+    data: Dict[float, List[str]] = {}
     for allocator in bench.results["allocators"]:
         value = _eval_with_stat(bench, datapoint, allocator, perm, "mean")
         if value in data: