"""
Base engine mixin — ``__init__``, scaling, placeholder, and conversion helpers.
"""
import logging
import io
import base64
from pathlib import Path
from typing import Dict, List, Any, Optional, Union
# Core plotting libraries
try:
import matplotlib.pyplot as plt
import seaborn as sns
MATPLOTLIB_AVAILABLE = True
except ImportError:
MATPLOTLIB_AVAILABLE = False
plt = None
sns = None
# Geographic plotting
try:
import folium
FOLIUM_AVAILABLE = True
except ImportError:
FOLIUM_AVAILABLE = False
folium = None
# Data processing
try:
import pandas as pd
import numpy as np
PANDAS_AVAILABLE = True
except ImportError:
PANDAS_AVAILABLE = False
pd = None
np = None
try:
from reportlab.platypus import Image, Paragraph, Spacer
from reportlab.lib.units import inch
from reportlab.lib.styles import getSampleStyleSheet
REPORTLAB_AVAILABLE = True
except ImportError:
REPORTLAB_AVAILABLE = False
Image = Paragraph = Spacer = None
inch = None
getSampleStyleSheet = None
log = logging.getLogger(__name__)
def _require_reportlab() -> None:
if not REPORTLAB_AVAILABLE:
raise ImportError(
"reportlab is required for PDF chart generation but is not installed. "
"Install it with: pip install reportlab"
)
def _rl_image_from_data_uri(data_uri: str, width_pts: float, height_pts: float) -> "Image":
"""Build a ReportLab Image from a base64 ``data:`` URI.
ReportLab 5 cannot construct an Image directly from a data: URI (its
open_for_read gates the data scheme behind trustedHosts and raises at
construction). Decode the URI to an in-memory buffer for rendering, then
keep the URI on ``.filename`` so image_utils can still extract the PNG.
"""
b64 = data_uri.split(",", 1)[1]
buffer = io.BytesIO(base64.b64decode(b64))
image = Image(buffer, width=width_pts, height=height_pts)
image.filename = data_uri
return image
__all__ = [
"BaseChartEngine",
]
[docs]
class BaseChartEngine:
"""Core infrastructure: init, colors, styling, scaling, save helpers."""
[docs]
def __init__(self, branding_config: Optional[Dict[str, Any]] = None,
output_dir: Optional[Path] = None,
max_chart_width: Optional[float] = None,
max_chart_height: Optional[float] = None):
"""
Initialize the chart generator.
Args:
branding_config: Branding configuration for chart colors and styling
output_dir: Directory for saving Folium HTML map files.
Defaults to ``~/.siege_utilities`` for backward compatibility.
max_chart_width: Maximum chart width in inches for ReportLab output.
When set, charts wider than this are uniformly scaled
down to fit. ``None`` disables width clamping.
max_chart_height: Maximum chart height in inches for ReportLab output.
When set, charts taller than this are uniformly scaled
down to fit. ``None`` disables height clamping.
"""
self.branding_config = branding_config or {}
self.output_dir = Path(output_dir) if output_dir is not None else Path.home() / ".siege_utilities"
self.max_chart_width = max_chart_width
self.max_chart_height = max_chart_height
self.styles = getSampleStyleSheet() if REPORTLAB_AVAILABLE else None
self._setup_default_colors()
self._setup_plotting_style()
# Set default figure size and DPI for professional quality
self.default_figsize = (8, 5)
self.default_dpi = 150
# Set matplotlib global DPI for consistency
if MATPLOTLIB_AVAILABLE:
plt.rcParams['figure.dpi'] = 150
plt.rcParams['savefig.dpi'] = 150
def _setup_default_colors(self):
"""Setup default color scheme for charts."""
self.default_colors = {
'primary': self.branding_config.get('colors', {}).get('primary', '#0077CC'),
'secondary': self.branding_config.get('colors', {}).get('secondary', '#EEEEEE'),
'accent': self.branding_config.get('colors', {}).get('accent', '#FF6B35'),
'success': '#28a745',
'warning': '#ffc107',
'danger': '#dc3545',
'info': '#17a2b8',
'light': '#f8f9fa',
'dark': '#343a40'
}
# Color palette for multiple series
self.color_palette = [
self.default_colors['primary'],
self.default_colors['accent'],
self.default_colors['success'],
self.default_colors['warning'],
self.default_colors['info']
]
def _setup_plotting_style(self):
"""Setup matplotlib and seaborn styling."""
if MATPLOTLIB_AVAILABLE:
# Set style
plt.style.use('default')
# Configure seaborn
if sns:
sns.set_palette(self.color_palette)
sns.set_style("whitegrid")
# Set default font sizes and DPI
plt.rcParams['figure.dpi'] = 150
plt.rcParams['font.size'] = 10
plt.rcParams['axes.titlesize'] = 12
plt.rcParams['axes.labelsize'] = 10
plt.rcParams['xtick.labelsize'] = 9
plt.rcParams['ytick.labelsize'] = 9
plt.rcParams['legend.fontsize'] = 9
plt.rcParams['figure.titlesize'] = 14
# Enhanced legend settings to prevent label collisions
plt.rcParams['legend.frameon'] = True
plt.rcParams['legend.fancybox'] = True
plt.rcParams['legend.shadow'] = True
plt.rcParams['legend.framealpha'] = 0.9
plt.rcParams['legend.edgecolor'] = 'black'
plt.rcParams['legend.facecolor'] = 'white'
def _scale_dimensions(self, width: float, height: float,
max_width: Optional[float] = None,
max_height: Optional[float] = None) -> tuple:
"""Return ``(width, height)`` uniformly scaled to fit within max bounds.
Resolves effective maximums from instance-level defaults
(``self.max_chart_width``/``self.max_chart_height``) and per-call
overrides. When both are set the *tighter* (smaller) limit wins.
If no maximum is active, the original dimensions are returned
unchanged. Aspect ratio is always preserved.
Args:
width: Original width in inches.
height: Original height in inches.
max_width: Per-call maximum width override (``None`` → use instance default).
max_height: Per-call maximum height override (``None`` → use instance default).
Returns:
``(scaled_width, scaled_height)`` tuple.
"""
# Resolve effective max — tighter of instance-level and per-call
eff_max_w = self.max_chart_width
if max_width is not None:
eff_max_w = min(max_width, eff_max_w) if eff_max_w is not None else max_width
eff_max_h = self.max_chart_height
if max_height is not None:
eff_max_h = min(max_height, eff_max_h) if eff_max_h is not None else max_height
# No limits → nothing to do
if eff_max_w is None and eff_max_h is None:
return width, height
scale = 1.0
if eff_max_w is not None and width > eff_max_w:
scale = min(scale, eff_max_w / width)
if eff_max_h is not None and height > eff_max_h:
scale = min(scale, eff_max_h / height)
if scale < 1.0:
new_w = width * scale
new_h = height * scale
log.info(
"Auto-scaling chart from %.2f×%.2f to %.2f×%.2f in "
"(max %.1s×%.1s, scale=%.3f)",
width, height, new_w, new_h,
eff_max_w, eff_max_h, scale,
)
return new_w, new_h
return width, height
def _save_folium_map(self, folium_map, filename: str, label: str,
width: float, height: float) -> Image:
"""Save a Folium map to *self.output_dir* and return a placeholder chart.
Consolidates the identical mkdir → save → log → placeholder pattern
used by every Folium-based method.
Args:
folium_map: A ``folium.Map`` instance.
filename: File name (e.g. ``"temp_choropleth_map.html"``).
label: Human-readable label shown in the placeholder image.
width: Placeholder chart width (inches).
height: Placeholder chart height (inches).
Returns:
ReportLab Image placeholder.
"""
map_path = self.output_dir / filename
map_path.parent.mkdir(parents=True, exist_ok=True)
folium_map.save(str(map_path))
log.info(f"{label} saved to {map_path}")
return self._create_placeholder_chart(width, height, f"{label} — saved to {map_path}")
def _matplotlib_to_reportlab_image(self, fig, width: float, height: float,
max_width: Optional[float] = None,
max_height: Optional[float] = None,
vector_export_path: Optional[Union[str, Path]] = None,
vector_format: str = 'svg') -> Image:
"""
Convert matplotlib figure to ReportLab Image with size optimization.
Optionally saves a vector copy (SVG/EPS/PDF) alongside the raster
PDF embed. The vector file is intended for designer handoff —
InDesign and Illustrator can import SVG/EPS directly.
Args:
fig: Matplotlib figure
width: Image width in inches
height: Image height in inches
max_width: Per-call max width override for ``_scale_dimensions``.
max_height: Per-call max height override for ``_scale_dimensions``.
vector_export_path: If provided, save a vector copy to this path.
vector_format: Format for vector export (``'svg'``, ``'eps'``, ``'pdf'``).
Returns:
ReportLab Image object
"""
_require_reportlab()
try:
# Save vector copy if requested (before closing the figure)
if vector_export_path is not None:
self.save_figure_as_vector(fig, vector_export_path, fmt=vector_format)
# Use high DPI for crisp, professional quality
optimal_dpi = self.default_dpi
# Save figure to bytes buffer with optimized settings
img_buffer = io.BytesIO()
fig.savefig(img_buffer, format='png', dpi=optimal_dpi,
facecolor='white', edgecolor='none', pad_inches=0.1)
img_buffer.seek(0)
# Check if image is still too large
img_size = len(img_buffer.getvalue())
if img_size > 5 * 1024 * 1024: # 5MB limit
log.warning(f"Image size {img_size/1024/1024:.1f}MB exceeds limit, reducing quality")
img_buffer.seek(0)
fig.savefig(img_buffer, format='png', dpi=100,
facecolor='white', edgecolor='none', pad_inches=0.1)
img_buffer.seek(0)
# Encode as base64
img_data = base64.b64encode(img_buffer.getvalue()).decode()
# Create data URL
data_url = f"data:image/png;base64,{img_data}"
# Close the figure to free memory
plt.close(fig)
# Auto-scale to fit within max bounds (if configured)
scaled_w, scaled_h = self._scale_dimensions(
width, height, max_width=max_width, max_height=max_height,
)
# ReportLab 5 cannot construct an Image from a data: URI -- its
# open_for_read gates the data scheme behind trustedHosts and fails
# at construction. Build the Image from the in-memory PNG buffer
# (which renders correctly), then keep the data URI on .filename so
# image_utils can still extract/save the embedded PNG. Width and
# height are set explicitly so ReportLab does not size from raw
# pixels (800px would otherwise become 800pt = 11.1in).
img_buffer.seek(0)
rl_image = Image(img_buffer, width=scaled_w * inch,
height=scaled_h * inch)
rl_image.filename = data_url
return rl_image
except (ValueError, TypeError, KeyError, IndexError, AttributeError, OSError) as e:
raise RuntimeError(
f"Image Conversion Error: {e}"
) from e
def _create_placeholder_chart(self, width: float, height: float,
text: str = "Chart Placeholder") -> Image:
"""
Create a placeholder chart when chart generation fails.
Args:
width: Chart width in inches
height: Chart height in inches
text: Text to display in placeholder
Returns:
ReportLab Image object
"""
_require_reportlab()
try:
# Create a simple placeholder using matplotlib
if MATPLOTLIB_AVAILABLE:
fig, ax = plt.subplots(figsize=(width, height), dpi=self.default_dpi)
ax.text(0.5, 0.5, text, ha='center', va='center', transform=ax.transAxes,
fontsize=12, bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgray"))
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis('off')
# Convert to ReportLab Image (inherits auto-scaling)
return self._matplotlib_to_reportlab_image(fig, width, height)
else:
# Fallback to simple text — still honour auto-scaling
sw, sh = self._scale_dimensions(width, height)
return _rl_image_from_data_uri(
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==",
sw * inch, sh * inch,
)
except (ValueError, TypeError, AttributeError, OSError) as e:
log.error(f"Error creating placeholder chart: {e}")
# Return a minimal image — still honour auto-scaling
sw, sh = self._scale_dimensions(width, height)
return _rl_image_from_data_uri(
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==",
sw * inch, sh * inch,
)
[docs]
def create_chart_with_caption(self, chart: Image, caption: str,
caption_style: Optional[str] = None) -> List:
"""
Create a chart with a caption below it.
Args:
chart: The chart image
caption: Caption text
caption_style: Style name for the caption
Returns:
List containing chart and caption
"""
_require_reportlab()
if caption_style is None:
caption_style = 'Caption' if 'Caption' in self.styles else 'Normal'
return [chart, Paragraph(caption, self.styles[caption_style]), Spacer(1, 0.2 * inch)]
[docs]
def create_chart_section(self, title: str, charts: List[Image],
description: str = "", width: float = 6.0) -> List:
"""
Create a complete chart section with title, description, and charts.
Args:
title: Section title
charts: List of chart images
description: Section description
width: Maximum section width in inches. Charts exceeding this
width are scaled down proportionally.
Returns:
List of flowables for the section
"""
_require_reportlab()
section = []
# Add title
if title:
section.append(Paragraph(title, self.styles['h2']))
section.append(Spacer(1, 0.1 * inch))
# Add description
if description:
section.append(Paragraph(description, self.styles['BodyText']))
section.append(Spacer(1, 0.2 * inch))
# Add charts, scaling to fit width constraint
max_width = width * inch
for chart in charts:
if hasattr(chart, 'drawWidth') and chart.drawWidth > max_width:
scale = max_width / chart.drawWidth
chart.drawWidth = max_width
chart.drawHeight = chart.drawHeight * scale
section.append(chart)
section.append(Spacer(1, 0.2 * inch))
return section
[docs]
def generate_chart_from_dataframe(self, df, chart_type: str = "bar",
x_column: str = None, y_columns: List[str] = None,
title: str = "", width: float = 6.0,
height: float = 4.0) -> Image:
"""
Generate a chart from a pandas DataFrame.
Args:
df: Pandas DataFrame
chart_type: Type of chart to create
x_column: Column to use for X-axis labels
y_columns: Columns to use for Y-axis data
title: Chart title
width: Chart width in inches
height: Chart height in inches
Returns:
ReportLab Image object
"""
try:
if chart_type == "bar":
return self.create_bar_chart(df, x_column, y_columns[0] if y_columns else None, title, width, height)
elif chart_type == "line":
return self.create_line_chart(df, x_column, y_columns, title, width, height)
elif chart_type == "pie":
return self.create_pie_chart(df, x_column, y_columns[0] if y_columns else None, title, width, height)
elif chart_type == "scatter":
return self.create_scatter_plot(df, x_column, y_columns[0] if y_columns else None, title, width, height)
elif chart_type == "heatmap":
return self.create_heatmap(df, title=title, width=width, height=height)
elif chart_type == "choropleth":
return self.create_choropleth_map(df, x_column, y_columns[0] if y_columns else None, title, width, height)
elif chart_type == "bivariate_choropleth":
if y_columns and len(y_columns) >= 2:
return self.create_bivariate_choropleth(df, x_column, y_columns[0], y_columns[1], title, width, height)
else:
raise ValueError("Bivariate choropleth requires at least 2 value columns")
else:
return self.create_bar_chart(df, x_column, y_columns[0] if y_columns else None, title, width, height)
except (ValueError, TypeError, KeyError, IndexError, AttributeError) as e:
raise RuntimeError(
f"DataFrame Chart Error: {e}"
) from e
[docs]
def create_custom_chart(self, chart_config: Dict[str, Any],
width: float = 6.0, height: float = 4.0) -> Image:
"""
Create a custom chart based on configuration.
Args:
chart_config: Complete chart configuration
width: Chart width in inches
height: Chart height in inches
Returns:
ReportLab Image object
"""
try:
# Validate chart configuration
required_keys = ['type', 'data']
if not all(key in chart_config for key in required_keys):
raise ValueError(f"Chart configuration missing required keys: {required_keys}")
chart_type = chart_config['type']
if chart_type == 'bar':
return self.create_bar_chart(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
elif chart_type == 'line':
return self.create_line_chart(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
elif chart_type == 'pie':
return self.create_pie_chart(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
elif chart_type == 'scatter':
return self.create_scatter_plot(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
elif chart_type == 'heatmap':
return self.create_heatmap(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
elif chart_type == 'choropleth':
return self.create_choropleth_map(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
elif chart_type == 'bivariate_choropleth':
# For bivariate choropleth, we need additional configuration
if 'location_column' in chart_config and 'value_columns' in chart_config and len(chart_config['value_columns']) >= 2:
return self.create_bivariate_choropleth(
chart_config['data'],
chart_config['location_column'],
chart_config['value_columns'][0],
chart_config['value_columns'][1],
title=chart_config.get('title', ''),
width=width,
height=height
)
else:
raise ValueError("Bivariate choropleth requires location_column and value_columns configuration")
else:
return self.create_bar_chart(chart_config['data'], title=chart_config.get('title', ''), width=width, height=height)
except (ValueError, TypeError, KeyError, IndexError, AttributeError) as e:
raise RuntimeError(
f"Custom Chart Error: {e}"
) from e