diff options
| -rw-r--r-- | allocbench/plots.py | 12 |
1 files changed, 7 insertions, 5 deletions
diff --git a/allocbench/plots.py b/allocbench/plots.py index a3cea75..c5340cb 100644 --- a/allocbench/plots.py +++ b/allocbench/plots.py @@ -25,7 +25,7 @@ import os import re import traceback from typing import (Any, cast, Mapping, MutableMapping, Dict, List, Tuple, - NamedTuple, Union, Sequence) + NamedTuple, Union, Sequence, Literal) import matplotlib import matplotlib.pyplot as plt @@ -47,6 +47,8 @@ SUMMARY_FILE_EXT = "svg" LATEX_CUSTOM_PREAMBLE = "" +PlotType = Literal['plot', 'errorbar', 'bar'] + DEFAULT_PLOT_OPTIONS = { 'plot': { 'marker': '.', @@ -184,7 +186,7 @@ def get_y_data(bench: Benchmark, return y_data -def _create_plot_options(plot_type: str, **kwargs) -> Dict[str, Any]: +def _create_plot_options(plot_type: PlotType, **kwargs) -> Dict[str, Any]: """ Create a plot options dictionary. @@ -212,7 +214,7 @@ def _create_plot_options(plot_type: str, **kwargs) -> Dict[str, Any]: return options -def _create_figure_options(plot_type: str, fig_label: str, +def _create_figure_options(plot_type: PlotType, fig_label: str, **kwargs) -> Dict[str, Any]: """ Create a figure options dictionary @@ -248,7 +250,7 @@ def _plot(bench: Benchmark, y_expression: str, x_data: Sequence, perms: Union[List[NamedTuple], NamedTuple], - plot_type: str, + plot_type: PlotType, plot_options: MutableMapping[str, Any], fig_options: Mapping[str, Any], scale: str = None, @@ -333,7 +335,7 @@ def _plot(bench: Benchmark, def plot(bench, y_expression, - plot_type='errorbar', + plot_type: PlotType = 'errorbar', x_args=None, scale=None, plot_options=None, |