From 8bf37e72d41360a920db0df78c6f389f41c879dd Mon Sep 17 00:00:00 2001 From: Florian Fischer Date: Mon, 20 Jul 2020 12:03:51 +0200 Subject: [plots] use Literal to construct PlotType --- allocbench/plots.py | 12 +++++++----- 1 file 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, -- cgit v1.2.3