[benchmark.py] add function to retrieve Benchmark instance from a module

3 files changed, 30 insertions, 21 deletions

diff --git a/allocbench/benchmark.py b/allocbench/benchmark.py
index c4b3695..642737a 100644
--- a/allocbench/benchmark.py
+++ b/allocbench/benchmark.py
@@ -21,6 +21,7 @@ from collections import namedtuple
 import errno
 import copy
 import csv
+import importlib
 import itertools
 import json
 import multiprocessing
@@ -726,3 +727,16 @@ class Benchmark:
                     stats["outliers"][key] = outliers
 
                 self.results["stats"][alloc][perm] = stats
+
+
+def get_benchmark_object(benchmark_name: str) -> 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}")
+    # find Benchmark class
+    for member in bench_module.__dict__.values():
+        if (not isinstance(member, type) or member is Benchmark
+                or not issubclass(member, Benchmark)):
+            continue
+
+        return member()
diff --git a/bench.py b/bench.py
index 007dbe9..adecc89 100755
--- a/bench.py
+++ b/bench.py
@@ -21,13 +21,13 @@
 import argparse
 import atexit
 import datetime
-import importlib
 import os
 import sys
 import traceback
 
 from allocbench.allocator import collect_allocators
 from allocbench.analyse import analyze_bench, analyze_allocators
+from allocbench.benchmark import get_benchmark_object
 import allocbench.facter as facter
 import allocbench.globalvars
 from allocbench.util import run_cmd
@@ -173,15 +173,13 @@ def main():
         if args.exclude_benchmarks and bench in args.exclude_benchmarks:
             continue
 
-        bench_module = importlib.import_module(
-            f"allocbench.benchmarks.{bench}")
-
-        # find Benchmark class
-        for member in bench_module.__dict__.values():
-            if (not isinstance(member, type)
-                    or member is allocbench.benchmark.Benchmark
-                    or not issubclass(member, allocbench.benchmark.Benchmark)):
-                continue
+        try:
+            print_status("Loading", bench, "...")
+            bench = get_benchmark_object(bench)
+        except Exception:
+            print_error(traceback.format_exc())
+            print_error(f"Skipping {bench}! Loading failed.")
+            continue
 
             try:
                 print_status("Preparing", bench, "...")
diff --git a/summarize.py b/summarize.py
index 9a37789..95a2ae4 100755
--- a/summarize.py
+++ b/summarize.py
@@ -19,12 +19,12 @@
 """Summarize the results of an allocbench run"""
 
 import argparse
-import importlib
 import os
 import sys
 
 import allocbench.facter as facter
 import allocbench.globalvars
+import allocbench.benchmark
 from allocbench.util import print_status, print_debug, print_error
 from allocbench.util import print_license_and_exit
 
@@ -72,7 +72,7 @@ def bench_sum(bench, exclude_allocators=None, sets=False):
     new_allocs = {
         an: a
         for an, a in bench.results["allocators"].items()
-        if an not in exclude_allocators or {}
+        if an not in (exclude_allocators or {})
     }
     bench.results["allocators"] = new_allocs
 
@@ -127,14 +127,11 @@ def summarize(benchmarks=None,
         if exclude_benchmarks and benchmark in exclude_benchmarks:
             continue
 
-        bench_module = importlib.import_module(
-            f"allocbench.benchmarks.{benchmark}")
-
-        if not hasattr(bench_module, benchmark):
-            print_error(f"{benchmark} has no member {benchmark}")
-            print_error(f"Skipping {benchmark}.")
-
-        bench = getattr(bench_module, benchmark, sets=sets)
+        try:
+            bench = allocbench.benchmark.get_benchmark_object(benchmark)
+        except Exception:  #pylint: disable=broad-except
+            print_error(f"Skipping {benchmark}. Loading failed")
+            continue
 
         try:
             bench.load()
@@ -143,7 +140,7 @@ def summarize(benchmarks=None,
 
         print_status(f"Summarizing {bench.name} ...")
         try:
-            bench_sum(bench, exclude_allocators=exclude_allocators)
+            bench_sum(bench, exclude_allocators=exclude_allocators, sets=sets)
         except FileExistsError as err:
             print(err)