|
|
import matplotlib.pyplot as plt |
|
|
import pandas as pd |
|
|
import os |
|
|
import logging |
|
|
import time |
|
|
|
|
|
class ChartGenerator: |
|
|
def __init__(self, data=None): |
|
|
logging.info("Initializing ChartGenerator") |
|
|
if data is not None: |
|
|
self.data = data |
|
|
else: |
|
|
self.data = pd.read_excel(os.path.join(os.path.dirname(os.path.dirname(__file__)), 'data', 'sample_data.xlsx')) |
|
|
|
|
|
def generate_chart(self, plot_args): |
|
|
start_time = time.time() |
|
|
logging.info(f"Generating chart with arguments: {plot_args}") |
|
|
|
|
|
|
|
|
x_col = plot_args['x'] |
|
|
y_cols = plot_args['y'] |
|
|
missing_cols = [] |
|
|
if x_col not in self.data.columns: |
|
|
missing_cols.append(x_col) |
|
|
for y in y_cols: |
|
|
if y not in self.data.columns: |
|
|
missing_cols.append(y) |
|
|
if missing_cols: |
|
|
logging.error(f"Missing columns in data: {missing_cols}") |
|
|
raise ValueError(f"Missing columns in data: {missing_cols}") |
|
|
|
|
|
fig, ax = plt.subplots() |
|
|
for y in y_cols: |
|
|
color = plot_args.get('color', None) |
|
|
if plot_args.get('chart_type', 'line') == 'bar': |
|
|
ax.bar(self.data[x_col], self.data[y], label=y, color=color) |
|
|
else: |
|
|
ax.plot(self.data[x_col], self.data[y], label=y, color=color) |
|
|
|
|
|
ax.set_xlabel(x_col) |
|
|
ax.legend() |
|
|
|
|
|
|
|
|
chart_filename = 'chart.png' |
|
|
output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'static', 'images') |
|
|
if not os.path.exists(output_dir): |
|
|
os.makedirs(output_dir) |
|
|
|
|
|
full_path = os.path.join(output_dir, chart_filename) |
|
|
|
|
|
if os.path.exists(full_path): |
|
|
os.remove(full_path) |
|
|
|
|
|
plt.savefig(full_path) |
|
|
|
|
|
logging.info(f"Chart generated and saved to {full_path}") |
|
|
|
|
|
return os.path.join('static', 'images', chart_filename) |
