use std::cell::RefCell;
use std::path::Path;
use crate::core::color::Rgba;
use crate::core::rcparams::RcParams;
use crate::figure::{Axes, Figure};
use crate::skia::PngError;
const DEFAULT_WIDTH_IN: f64 = 6.4;
const DEFAULT_HEIGHT_IN: f64 = 4.8;
const DEFAULT_AXES_RECT: (f64, f64, f64, f64) = (0.125, 0.11, 0.775, 0.79);
fn new_single_axes_figure(w_in: f64, h_in: f64, rc: &RcParams) -> Figure {
let mut fig = Figure::new(w_in, h_in).with_rcparams(rc.clone());
let (l, b, w, h) = DEFAULT_AXES_RECT;
fig.add_axes(l, b, w, h);
fig
}
struct PyplotState {
figure: Option<Figure>,
current_axes: usize,
style: Option<RcParams>,
}
impl PyplotState {
const fn new() -> Self {
Self {
figure: None,
current_axes: 0,
style: None,
}
}
fn current_style(&self) -> RcParams {
self.style.clone().unwrap_or_default()
}
fn ensure_figure(&mut self) -> &mut Figure {
if self.figure.is_none() {
let rc = self.current_style();
self.figure = Some(new_single_axes_figure(
DEFAULT_WIDTH_IN,
DEFAULT_HEIGHT_IN,
&rc,
));
self.current_axes = 0;
}
self.figure.as_mut().expect("figure just ensured")
}
fn ensure_axes(&mut self) -> &mut Axes {
let index = self.current_axes;
let fig = self.ensure_figure();
fig.axes_mut()
.get_mut(index)
.expect("current axes index in range")
}
}
thread_local! {
static STATE: RefCell<PyplotState> = const { RefCell::new(PyplotState::new()) };
}
fn with_axes<R>(f: impl FnOnce(&mut Axes) -> R) -> R {
STATE.with(|s| f(s.borrow_mut().ensure_axes()))
}
fn with_figure<R>(f: impl FnOnce(&mut Figure) -> R) -> R {
STATE.with(|s| f(s.borrow_mut().ensure_figure()))
}
pub fn figure() {
figure_sized(DEFAULT_WIDTH_IN, DEFAULT_HEIGHT_IN);
}
pub fn figure_sized(w_in: f64, h_in: f64) {
STATE.with(|s| {
let mut state = s.borrow_mut();
let rc = state.current_style();
state.figure = Some(new_single_axes_figure(w_in, h_in, &rc));
state.current_axes = 0;
});
}
pub fn style(rc: RcParams) {
STATE.with(|s| s.borrow_mut().style = Some(rc));
}
pub fn subplots(nrows: usize, ncols: usize) {
assert!(nrows >= 1 && ncols >= 1, "subplots grid must be non-empty");
STATE.with(|s| {
let mut state = s.borrow_mut();
let rc = state.current_style();
let mut fig = Figure::new(DEFAULT_WIDTH_IN, DEFAULT_HEIGHT_IN).with_rcparams(rc);
for index in 1..=(nrows * ncols) {
fig.add_subplot(nrows, ncols, index);
}
state.figure = Some(fig);
state.current_axes = 0;
});
}
pub fn sca(index: usize) {
STATE.with(|s| {
let mut state = s.borrow_mut();
let len = state.ensure_figure().axes().len();
if index < len {
state.current_axes = index;
}
});
}
pub fn plot(x: &[f64], y: &[f64]) {
with_axes(|ax| {
ax.plot(x, y);
});
}
pub fn scatter(x: &[f64], y: &[f64]) {
with_axes(|ax| {
ax.scatter(x, y);
});
}
pub fn bar(x: &[f64], height: &[f64]) {
with_axes(|ax| ax.bar(x, height));
}
pub fn hist(data: &[f64], bins: usize) {
with_axes(|ax| {
ax.hist(data, bins);
});
}
pub fn fill_between(x: &[f64], y1: &[f64], y2: &[f64]) {
with_axes(|ax| ax.fill_between(x, y1, y2));
}
pub fn barh(y: &[f64], width: &[f64]) {
with_axes(|ax| ax.barh(y, width));
}
pub fn step(x: &[f64], y: &[f64]) {
with_axes(|ax| {
ax.step(x, y);
});
}
pub fn stairs(values: &[f64], edges: &[f64]) {
with_axes(|ax| ax.stairs(values, edges));
}
pub fn stem(x: &[f64], y: &[f64]) {
with_axes(|ax| ax.stem(x, y));
}
pub fn stackplot(x: &[f64], ys: &[&[f64]]) {
with_axes(|ax| ax.stackplot(x, ys));
}
pub fn broken_barh(xranges: &[(f64, f64)], yrange: (f64, f64)) {
with_axes(|ax| ax.broken_barh(xranges, yrange));
}
pub fn grouped_bar(series: &[&[f64]]) {
with_axes(|ax| {
ax.grouped_bar(series);
});
}
pub fn errorbar(x: &[f64], y: &[f64], yerr: &[f64]) {
with_axes(|ax| ax.errorbar(x, y, yerr));
}
pub fn boxplot(data: &[&[f64]]) {
with_axes(|ax| ax.boxplot(data));
}
pub fn violinplot(data: &[&[f64]], positions: Option<&[f64]>) {
with_axes(|ax| {
ax.violinplot(data, positions);
});
}
pub fn hexbin(x: &[f64], y: &[f64], gridsize: usize) {
with_axes(|ax| {
ax.hexbin(x, y, gridsize);
});
}
pub fn ecdf(data: &[f64]) {
with_axes(|ax| ax.ecdf(data));
}
pub fn imshow(data: &[f64], nrows: usize, ncols: usize) {
with_axes(|ax| {
ax.imshow(data, nrows, ncols);
});
}
pub fn matshow(data: &[f64], nrows: usize, ncols: usize) {
with_axes(|ax| {
ax.matshow(data, nrows, ncols);
});
}
pub fn spy(data: &[f64], nrows: usize, ncols: usize) {
with_axes(|ax| {
ax.spy(data, nrows, ncols);
});
}
pub fn pcolormesh(c: &[f64], nrows: usize, ncols: usize) {
with_axes(|ax| {
ax.pcolormesh(c, nrows, ncols);
});
}
pub fn contour(z: &[f64], nrows: usize, ncols: usize) {
with_axes(|ax| ax.contour(z, nrows, ncols));
}
pub fn contour_levels(z: &[f64], nrows: usize, ncols: usize, n_levels: usize) {
with_axes(|ax| ax.contour_levels(z, nrows, ncols, n_levels));
}
pub fn hist2d(x: &[f64], y: &[f64], bins: usize) {
with_axes(|ax| {
ax.hist2d(x, y, bins);
});
}
pub fn eventplot(positions: &[&[f64]]) {
with_axes(|ax| ax.eventplot(positions));
}
pub fn fill_betweenx(y: &[f64], x1: &[f64], x2: &[f64]) {
with_axes(|ax| ax.fill_betweenx(y, x1, x2));
}
pub fn pie(values: &[f64]) {
with_axes(|ax| {
ax.pie(values);
});
}
pub fn hlines(y: &[f64], xmin: f64, xmax: f64) {
with_axes(|ax| ax.hlines(y, xmin, xmax));
}
pub fn vlines(x: &[f64], ymin: f64, ymax: f64) {
with_axes(|ax| ax.vlines(x, ymin, ymax));
}
pub fn axhline(y: f64) {
with_axes(|ax| ax.axhline(y));
}
pub fn axvline(x: f64) {
with_axes(|ax| ax.axvline(x));
}
pub fn axhspan(ymin: f64, ymax: f64) {
with_axes(|ax| ax.axhspan(ymin, ymax));
}
pub fn axvspan(xmin: f64, xmax: f64) {
with_axes(|ax| ax.axvspan(xmin, xmax));
}
pub fn axis_equal() {
with_axes(|ax| {
ax.set_aspect_equal();
});
}
pub fn axis_off() {
with_axes(|ax| {
ax.set_axis_off();
});
}
pub fn axis_on() {
with_axes(|ax| {
ax.set_axis_on();
});
}
pub fn margins(margin: f64) {
with_axes(|ax| {
ax.set_margins(margin);
});
}
pub fn legend(entries: Vec<(Rgba, String)>) {
with_axes(|ax| {
ax.legend(entries);
});
}
pub fn colorbar(cmap_name: &str, vmin: f64, vmax: f64) {
with_figure(|fig| {
fig.colorbar(cmap_name, vmin, vmax);
});
}
pub fn title(text: &str) {
with_axes(|ax| {
ax.set_title(text);
});
}
pub fn xlabel(text: &str) {
with_axes(|ax| {
ax.set_xlabel(text);
});
}
pub fn ylabel(text: &str) {
with_axes(|ax| {
ax.set_ylabel(text);
});
}
pub fn xlim(lo: f64, hi: f64) {
with_axes(|ax| {
ax.set_xlim(lo, hi);
});
}
pub fn ylim(lo: f64, hi: f64) {
with_axes(|ax| {
ax.set_ylim(lo, hi);
});
}
#[derive(Debug)]
pub enum SaveError {
Png(PngError),
Io(std::io::Error),
}
impl std::fmt::Display for SaveError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SaveError::Png(e) => write!(f, "PNG encoding failed: {e}"),
SaveError::Io(e) => write!(f, "writing figure file failed: {e}"),
}
}
}
impl std::error::Error for SaveError {}
pub fn savefig<P: AsRef<Path>>(path: P) -> Result<(), SaveError> {
let path = path.as_ref();
let ext = path
.extension()
.and_then(|e| e.to_str())
.map(str::to_ascii_lowercase);
with_figure(|fig| match ext.as_deref() {
Some("svg") => fig.save_svg(path).map_err(SaveError::Io),
Some("pdf") => fig.save_pdf(path).map_err(SaveError::Io),
_ => fig.save_png(path).map_err(SaveError::Png),
})
}
pub fn show() {
#[cfg(all(feature = "show", not(target_arch = "wasm32")))]
{
let figure = STATE.with(|s| {
let mut state = s.borrow_mut();
state.current_axes = 0;
state.figure.take()
});
if let Some(figure) = figure {
crate::show::show(figure);
}
}
#[cfg(not(all(feature = "show", not(target_arch = "wasm32"))))]
{
eprintln!("crate::pyplot::show() needs the `show` feature; use savefig() instead.");
}
}
pub fn clf() {
STATE.with(|s| {
let mut state = s.borrow_mut();
state.figure = None;
state.current_axes = 0;
});
}
pub fn close() {
clf();
}
#[cfg(test)]
mod tests {
use super::*;
fn target_path(name: &str) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
p.push(name);
p
}
#[test]
fn plot_title_savefig_writes_nonempty_png() {
figure();
let xs = [0.0, 1.0, 2.0, 3.0];
let ys = [0.0, 1.0, 0.0, 1.0];
plot(&xs, &ys);
title("t");
let path = target_path("pyplot_test.png");
savefig(&path).expect("savefig succeeds");
let meta = std::fs::metadata(&path).expect("file exists");
assert!(meta.len() > 0, "PNG should be non-empty");
close();
}
#[test]
fn savefig_picks_format_from_extension() {
figure();
plot(&[0.0, 1.0, 2.0], &[0.0, 1.0, 0.0]);
let svg = target_path("pyplot_ext.svg");
savefig(&svg).expect("svg savefig succeeds");
assert!(
std::fs::read_to_string(&svg)
.expect("svg file")
.contains("<svg"),
"expected an SVG document"
);
let pdf = target_path("pyplot_ext.pdf");
savefig(&pdf).expect("pdf savefig succeeds");
assert!(
std::fs::read(&pdf).expect("pdf file").starts_with(b"%PDF"),
"expected a PDF document"
);
let png = target_path("pyplot_ext.png");
savefig(&png).expect("png savefig succeeds");
assert!(
std::fs::read(&png)
.expect("png file")
.starts_with(&[0x89, b'P', b'N', b'G']),
"expected a PNG document"
);
close();
}
#[test]
fn plot_auto_initializes_without_figure() {
close();
plot(&[0.0, 1.0], &[1.0, 0.0]);
let path = target_path("pyplot_autoinit.png");
savefig(&path).expect("savefig succeeds after auto-init");
assert!(
std::fs::metadata(&path).expect("file exists").len() > 0,
"auto-init PNG should be non-empty"
);
close();
}
#[test]
fn close_resets_state_for_fresh_plot() {
plot(&[0.0, 5.0], &[0.0, 5.0]);
close();
STATE.with(|s| {
assert!(s.borrow().figure.is_none(), "close drops the figure");
});
plot(&[0.0, 1.0], &[0.0, 1.0]);
STATE.with(|s| {
let st = s.borrow();
let fig = st.figure.as_ref().expect("fresh figure exists");
assert_eq!(fig.axes().len(), 1, "fresh figure has one axes");
});
close();
}
#[test]
fn expanded_wrappers_smoke_savefig() {
figure();
let xs = [0.0, 1.0, 2.0, 3.0];
let ys = [1.0, 2.0, 1.5, 3.0];
barh(&xs, &ys);
step(&xs, &ys);
stairs(&[1.0, 2.0, 3.0], &[0.0, 1.0, 2.0, 3.0]);
errorbar(&xs, &ys, &[0.1, 0.2, 0.1, 0.3]);
let group_a = [1.0, 2.0, 3.0];
let group_b = [2.0, 1.0, 4.0];
let groups: [&[f64]; 2] = [&group_a, &group_b];
boxplot(&groups);
violinplot(&groups, None);
axhline(2.0);
axvline(1.5);
axhspan(0.5, 1.0);
margins(0.05);
pie(&[1.0, 2.0, 3.0]);
let path = target_path("pyplot_expanded_smoke.png");
savefig(&path).expect("savefig succeeds");
assert!(
std::fs::metadata(&path).expect("file exists").len() > 0,
"smoke PNG should be non-empty"
);
close();
}
#[test]
fn stackplot_threads_nested_slices() {
figure();
let x = [0.0, 1.0, 2.0];
let s0 = [1.0, 2.0, 1.0];
let s1 = [0.5, 0.5, 2.0];
let series: [&[f64]; 2] = [&s0, &s1];
stackplot(&x, &series);
let path = target_path("pyplot_stackplot.png");
savefig(&path).expect("savefig succeeds");
assert!(std::fs::metadata(&path).expect("file exists").len() > 0);
close();
}
#[test]
fn matrix_wrappers_accept_flat_data() {
figure();
let data = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0];
imshow(&data, 2, 3);
let path = target_path("pyplot_imshow.png");
savefig(&path).expect("savefig succeeds");
assert!(std::fs::metadata(&path).expect("file exists").len() > 0);
close();
}
#[test]
fn broken_barh_threads_tuple_args() {
figure();
broken_barh(&[(0.0, 1.0), (2.0, 1.5)], (1.0, 0.5));
let path = target_path("pyplot_broken_barh.png");
savefig(&path).expect("savefig succeeds");
assert!(std::fs::metadata(&path).expect("file exists").len() > 0);
close();
}
#[test]
fn subplots_sets_first_axes_current() {
subplots(2, 2);
STATE.with(|s| {
let st = s.borrow();
assert_eq!(st.figure.as_ref().expect("figure").axes().len(), 4);
assert_eq!(st.current_axes, 0);
});
close();
}
#[test]
fn legend_and_colorbar_render_to_png() {
figure();
plot(&[0.0, 1.0, 2.0], &[0.0, 1.0, 0.5]);
legend(vec![
(Rgba::new(0.2, 0.4, 0.8, 1.0), "series a".to_string()),
(Rgba::new(0.8, 0.3, 0.2, 1.0), "series b".to_string()),
]);
colorbar("viridis", 0.0, 1.0);
let path = target_path("pyplot_legend_colorbar.png");
savefig(&path).expect("savefig succeeds");
assert!(
std::fs::metadata(&path).expect("file exists").len() > 0,
"legend + colorbar PNG should be non-empty"
);
close();
}
}