"""Implementation of the LineGraphDrawer class to take given data and draw the graph."""
import sys
from copy import deepcopy
from pathlib import Path
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from jsonschema import exceptions, validate
from matplotlib import dates, ticker
from matplotlib.axes import Axes
from yaml import safe_load
from ...helpers.click_helpers import echo, recho
from .line_graph_drawer_interface import LineGraphDrawerInterface
from .settings_graph import Description, GraphSettings, LinesSettings, Mapping
class LineGraphDrawer(LineGraphDrawerInterface):
"""Class that implements the interface LineGraphDrawer"""
def __init__(
self,
graph: GraphSettings,
descriptions: Description,
mapping: Mapping,
axes: Axes,
name: str,
) -> None:
self._graph = graph
self._descriptions = descriptions
self._mapping = mapping
self._axes = axes
self.name = name
self._cycler = plt.rcParams["axes.prop_cycle"]()
def draw(self, data: pd.DataFrame) -> None:
"""Public method which draws the specified graph with error handling"""
try:
self._draw(data)
except KeyError as e:
recho(f"Error plot config file: Column {e!s} is not known")
sys.exit(1)
except ValueError as e:
recho(f"Some provided values can't be handled properly: {e!s}")
sys.exit(1)
except IndexError:
recho(
"Error plot config file: Number of y_axes descriptions/labels"
" does not match number of y-axes."
)
sys.exit(1)
except StopIteration:
recho(
"Error plot config file: Number of labels does not "
"match with number of lines."
)
sys.exit(1)
def show(self) -> None:
"""Shows the plot if wanted"""
if self._graph.show:
plt.show()
else:
plt.clf()
plt.close("all")
def save(self, output_dir: Path) -> None:
"""Saves the plot if wanted"""
if self._graph.save:
fname = Path(output_dir) / Path(self.name + "." + self._graph.format)
echo(f"Save plot at {fname}")
plt.savefig(
fname=fname,
dpi=self._graph.dpi,
format=self._graph.format,
)
def _draw(self, data: pd.DataFrame) -> None:
"""Private method which draws the specified graph, without
error handling
"""
x_values = data.loc[:, self._mapping.x].to_numpy()
if self._mapping.date_format is not None:
self._axes.xaxis.set_major_formatter(
dates.DateFormatter(self._mapping.date_format)
)
elif x_values.dtype.type is np.object_:
recho(
"The data type of the x values is not numerical or "
"datetime. The number of x-ticks can vary!",
"yellow",
)
plots: list[mpl.lines.Line2D] = []
for i, ax_name in enumerate(["y1", "y2", "y3"]):
line_settings = getattr(self._mapping, ax_name)
if isinstance(line_settings, LinesSettings):
if i >= 1:
axis = self._axes.twinx()
if ax_name == "y3":
axis.spines.right.set_position(("axes", 1.2))
plt.subplots_adjust(right=0.75)
else:
axis = self._axes
axis.set_ylabel(self._descriptions.y_axes[i])
for column in deepcopy(line_settings).input:
y_values = data.loc[:, column].to_numpy()
plots.extend(
self._draw_line(axis, x_values, y_values, line_settings)
)
self._axes.legend(handles=plots)
self._axes.set_title(self._descriptions.title)
self._axes.set_xlabel(self._descriptions.x_axis)
self._axes.xaxis.set_major_locator(
ticker.LinearLocator(self._mapping.x_ticks_count)
)
self._axes.margins(x=0)
def _draw_line(
self,
axis: Axes,
x_values: np.ndarray,
y_values: np.ndarray,
settings: LinesSettings,
) -> list[mpl.lines.Line2D]:
"""Draws the specified line"""
x_values = x_values[self._mapping.start : self._mapping.end]
y_values = y_values[self._mapping.start : self._mapping.end]
if not (y_values.dtype.type is np.str_ or y_values.dtype.type is np.object_):
x_values = x_values[~np.isnan(y_values)]
y_values = y_values[~np.isnan(y_values)]
if self._mapping.time_factor is not None:
x_values = x_values // self._mapping.time_factor
if self._mapping.start_date is not None:
deltas = x_values - x_values[0]
x_values = np.datetime64(self._mapping.start_date, "s") + deltas
scaley = True
if (settings.min is not None) and (settings.max is not None):
if y_values.dtype.type is np.str_ or y_values.dtype.type is np.object_:
recho(
"Min/Max axis limits for string y-values in plot "
f"'{self._descriptions.title}' are not allowed"
)
sys.exit(1)
axis.set_ylim(settings.min, settings.max)
scaley = False
if y_values.dtype.type is np.str_ or y_values.dtype.type is np.object_:
scaled_y_values = y_values
else:
scaled_y_values = y_values * settings.factor
label = next(settings.labels)
return axis.plot(
x_values,
scaled_y_values,
scaley=scaley,
label=label,
**next(self._cycler),
)
@staticmethod
def validate_config(config: dict) -> None:
"""Validates the CSVHandler configuration"""
schema_path = Path(__file__).parent / "schemas" / "line_graph_drawer.json"
with open(schema_path, encoding="utf-8") as f:
schema = safe_load(f)
try:
validate(config, schema=schema)
except exceptions.ValidationError as e:
error_text = str(e).splitlines()[0]
recho(f"Line graph config validation error: {error_text}")
sys.exit(1)