+class PerfRegressionDataPersister(ABC):
+ def __init__(self):
+ pass
+
+ @abstractmethod
+ def load_performance_data(self) -> Dict[str, List[float]]:
+ pass
+
+ @abstractmethod
+ def save_performance_data(self, method_id: str, data: Dict[str, List[float]]):
+ pass
+
+ @abstractmethod
+ def delete_performance_data(self, method_id: str):
+ pass
+
+
+class FileBasedPerfRegressionDataPersister(PerfRegressionDataPersister):
+ def __init__(self, filename: str):
+ self.filename = filename
+ self.traces_to_delete = []
+
+ def load_performance_data(self, method_id: str) -> Dict[str, List[float]]:
+ with open(self.filename, 'rb') as f:
+ return pickle.load(f)
+
+ def save_performance_data(self, method_id: str, data: Dict[str, List[float]]):
+ for trace in self.traces_to_delete:
+ if trace in data:
+ data[trace] = []
+
+ with open(self.filename, 'wb') as f:
+ pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
+
+ def delete_performance_data(self, method_id: str):
+ self.traces_to_delete.append(method_id)
+
+
+class DatabasePerfRegressionDataPersister(PerfRegressionDataPersister):
+ def __init__(self, dbspec: str):
+ self.dbspec = dbspec
+ self.engine = sa.create_engine(self.dbspec)
+ self.conn = self.engine.connect()
+
+ def load_performance_data(self, method_id: str) -> Dict[str, List[float]]:
+ results = self.conn.execute(
+ sa.text(
+ f'SELECT * FROM runtimes_by_function WHERE function = "{method_id}";'
+ )
+ )
+ ret = {method_id: []}
+ for result in results.all():
+ ret[method_id].append(result['runtime'])
+ results.close()
+ return ret
+
+ def save_performance_data(self, method_id: str, data: Dict[str, List[float]]):
+ self.delete_performance_data(method_id)
+ for (method_id, perf_data) in data.items():
+ sql = 'INSERT INTO runtimes_by_function (function, runtime) VALUES '
+ for perf in perf_data:
+ self.conn.execute(sql + f'("{method_id}", {perf});')
+
+ def delete_performance_data(self, method_id: str):
+ sql = f'DELETE FROM runtimes_by_function WHERE function = "{method_id}"'
+ self.conn.execute(sql)
+
+
+def function_identifier(f: Callable) -> str:
+ """
+ Given a callable function, return a string that identifies it.
+ Usually that string is just __module__:__name__ but there's a
+ corner case: when __module__ is __main__ (i.e. the callable is
+ defined in the same module as __main__). In this case,
+ f.__module__ returns "__main__" instead of the file that it is
+ defined in. Work around this using pathlib.Path (see below).
+
+ >>> function_identifier(function_identifier)
+ 'unittest_utils:function_identifier'
+
+ """
+ if f.__module__ == '__main__':
+ from pathlib import Path
+ import __main__
+ module = __main__.__file__
+ module = Path(module).stem
+ return f'{module}:{f.__name__}'
+ else:
+ return f'{f.__module__}:{f.__name__}'