use std::{
collections::HashMap,
fs,
io::{ self, Write },
ops::Range,
path::{ Path, PathBuf },
process,
rc::Rc,
};
use rand::distributions::{ Alphanumeric, DistString };
use serde_json as json;
use thiserror::Error;
use crate::commands::Axis2;
#[derive(Debug, Error)]
pub enum MplError {
#[error("IO error: {0}")]
IOError(#[from] io::Error),
#[error("serialization error: {0}")]
JsonError(#[from] json::Error),
#[error("script error:\nstdout:\n{0}\nstderr:\n{1}")]
PyError(String, String),
}
pub type MplResult<T> = Result<T, MplError>;
pub const PRELUDE: &str
= "\
import datetime
import io
import json
import os
import random
import sys
import matplotlib
matplotlib.use(\"QtAgg\")
import matplotlib.path as mpath
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import matplotlib.cm as mcm
import matplotlib.colors as mcolors
import matplotlib.collections as mcollections
import matplotlib.ticker as mticker
import matplotlib.image as mimage
from mpl_toolkits.mplot3d import axes3d
import numpy as np
";
pub const INIT: &str
= "\
fig, ax = plt.subplots()
";
pub trait Matplotlib: std::fmt::Debug {
fn is_prelude(&self) -> bool;
fn data(&self) -> Option<json::Value>;
fn py_cmd(&self) -> String;
}
pub trait AsPy {
fn as_py(&self) -> String;
}
impl AsPy for bool {
fn as_py(&self) -> String { if *self { "True" } else { "False" }.into() }
}
impl AsPy for i32 {
fn as_py(&self) -> String { self.to_string() }
}
impl AsPy for f64 {
fn as_py(&self) -> String { self.to_string() }
}
impl AsPy for String {
fn as_py(&self) -> String { format!("\"{self}\"") }
}
impl AsPy for &str {
fn as_py(&self) -> String { format!("\"{self}\"") }
}
#[derive(Clone, Debug, PartialEq)]
pub enum PyValue {
Bool(bool),
Int(i32),
Float(f64),
Str(String),
List(Vec<PyValue>),
Dict(HashMap<String, PyValue>),
Var(String),
None
}
impl From<bool> for PyValue {
fn from(b: bool) -> Self { Self::Bool(b) }
}
impl From<i32> for PyValue {
fn from(i: i32) -> Self { Self::Int(i) }
}
impl From<f64> for PyValue {
fn from(f: f64) -> Self { Self::Float(f) }
}
impl From<String> for PyValue {
fn from(s: String) -> Self { Self::Str(s) }
}
impl From<&str> for PyValue {
fn from(s: &str) -> Self { Self::Str(s.into()) }
}
impl<T> From<&T> for PyValue
where T: Clone + Into<PyValue>
{
fn from(x: &T) -> Self { x.clone().into() }
}
impl From<Vec<PyValue>> for PyValue {
fn from(l: Vec<PyValue>) -> Self { Self::List(l) }
}
impl From<HashMap<String, PyValue>> for PyValue {
fn from(d: HashMap<String, PyValue>) -> Self { Self::Dict(d) }
}
impl<T: Into<PyValue>> FromIterator<T> for PyValue {
fn from_iter<I>(iter: I) -> Self
where I: IntoIterator<Item = T>
{
Self::list(iter)
}
}
impl<S: Into<String>, T: Into<PyValue>> FromIterator<(S, T)> for PyValue {
fn from_iter<I>(iter: I) -> Self
where I: IntoIterator<Item = (S, T)>
{
Self::dict(iter)
}
}
impl PyValue {
pub fn list<I, T>(items: I) -> Self
where
I: IntoIterator<Item = T>,
T: Into<PyValue>,
{
Self::List(items.into_iter().map(|item| item.into()).collect())
}
pub fn dict<I, S, T>(items: I) -> Self
where
I: IntoIterator<Item = (S, T)>,
S: Into<String>,
T: Into<PyValue>,
{
Self::Dict(
items.into_iter().map(|(s, v)| (s.into(), v.into())).collect())
}
}
impl AsPy for PyValue {
fn as_py(&self) -> String {
match self {
Self::Bool(b) => if *b { "True".into() } else { "False".into() },
Self::Int(i) => format!("{i}"),
Self::Float(f) => format!("{f}"),
Self::Str(s) => format!("\"{s}\""),
Self::List(l) => {
let n = l.len();
let mut out = String::from("[");
for (k, v) in l.iter().enumerate() {
out.push_str(&v.as_py());
if k < n - 1 { out.push_str(", "); }
}
out.push(']');
out
},
Self::Dict(d) => {
let n = d.len();
let mut out = String::from("{");
for (j, (k, v)) in d.iter().enumerate() {
out.push_str(&format!("\"{}\": {}", k, v.as_py()));
if j < n - 1 { out.push_str(", "); }
}
out.push('}');
out
},
Self::Var(v) => v.clone(),
Self::None => "None".into(),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct Opt(pub String, pub PyValue);
impl<T: Into<PyValue>> From<(&str, T)> for Opt {
fn from(kv: (&str, T)) -> Self { Self(kv.0.into(), kv.1.into()) }
}
impl<T: Into<PyValue>> From<(String, T)> for Opt {
fn from(kv: (String, T)) -> Self { Self(kv.0, kv.1.into()) }
}
impl Opt {
pub fn new<T>(key: &str, val: T) -> Self
where T: Into<PyValue>
{
Self(key.into(), val.into())
}
}
pub fn opt<T>(key: &str, val: T) -> Opt
where T: Into<PyValue>
{
Opt::new(key, val)
}
impl AsPy for Opt {
fn as_py(&self) -> String { format!("{}={}", self.0, self.1.as_py()) }
}
impl AsPy for Vec<Opt> {
fn as_py(&self) -> String {
let n = self.len();
let mut out = String::new();
for (k, opt) in self.iter().enumerate() {
out.push_str(&opt.as_py());
if k < n - 1 { out.push_str(", "); }
}
out
}
}
pub trait MatplotlibOpts: Matplotlib {
fn kwarg<T: Into<PyValue>>(&mut self, key: &str, val: T) -> &mut Self;
fn o<T: Into<PyValue>>(mut self, key: &str, val: T) -> Self
where Self: Sized
{
self.kwarg(key, val);
self
}
fn oo<I>(mut self, opts: I) -> Self
where
I: IntoIterator<Item = Opt>,
Self: Sized,
{
opts.into_iter().for_each(|Opt(key, val)| { self.kwarg(&key, val); });
self
}
}
fn get_temp_fname() -> PathBuf {
std::env::temp_dir()
.join(Alphanumeric.sample_string(&mut rand::thread_rng(), 15))
}
#[derive(Debug)]
struct TempFile(PathBuf, Option<fs::File>);
impl TempFile {
fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
let path = path.as_ref().to_path_buf();
fs::OpenOptions::new()
.create(true)
.append(false)
.truncate(true)
.write(true)
.open(&path)
.map(|file| Self(path, Some(file)))
}
}
impl Write for TempFile {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if let Some(file) = self.1.as_mut() {
file.write(buf)
} else {
Ok(0)
}
}
fn flush(&mut self) -> io::Result<()> {
if let Some(file) = self.1.as_mut() {
file.flush()
} else {
Ok(())
}
}
}
impl Drop for TempFile {
fn drop(&mut self) {
drop(self.1.take());
fs::remove_file(&self.0).ok();
}
}
#[derive(Clone, Debug, Default)]
pub struct Mpl {
prelude: Vec<Rc<dyn Matplotlib + 'static>>,
commands: Vec<Rc<dyn Matplotlib + 'static>>,
}
impl Mpl {
pub fn new() -> Self { Self::default() }
pub fn new_3d<I>(opts: I) -> Self
where I: IntoIterator<Item = Opt>
{
let opts: Vec<Opt> = opts.into_iter().collect();
Self::default()
& crate::commands::DefPrelude
& crate::commands::Init3D { opts: opts.into_iter().collect() }
}
pub fn new_3d_with<I, F>(opts: I, f: F) -> Self
where
I: IntoIterator<Item = Opt>,
F: FnOnce(Mpl) -> Mpl,
{
f(Self::default() & crate::commands::DefPrelude)
& crate::commands::Init3D { opts: opts.into_iter().collect() }
}
pub fn new_grid<I>(nrows: usize, ncols: usize, opts: I) -> Self
where I: IntoIterator<Item = Opt>
{
let opts: Vec<Opt> = opts.into_iter().collect();
Self::default()
& crate::commands::DefPrelude
& crate::commands::InitGrid {
nrows,
ncols,
opts: opts.into_iter().collect(),
}
}
pub fn new_grid_with<I, F>(nrows: usize, ncols: usize, opts: I, f: F)
-> Self
where
I: IntoIterator<Item = Opt>,
F: FnOnce(Mpl) -> Mpl,
{
f(Self::default() & crate::commands::DefPrelude)
& crate::commands::InitGrid {
nrows,
ncols,
opts: opts.into_iter().collect(),
}
}
pub fn new_gridspec<I, P>(gridspec_kw: I, positions: P) -> Self
where
I: IntoIterator<Item = Opt>,
P: IntoIterator<Item = GSPos>,
{
Self::default()
& crate::commands::DefPrelude
& crate::commands::init_gridspec(gridspec_kw, positions)
}
pub fn new_gridspec_with<I, P, F>(gridspec_kw: I, positions: P, f: F)
-> Self
where
I: IntoIterator<Item = Opt>,
P: IntoIterator<Item = GSPos>,
F: FnOnce(Mpl) -> Mpl,
{
f(Self::default() & crate::commands::DefPrelude)
& crate::commands::init_gridspec(gridspec_kw, positions)
}
pub fn then<M: Matplotlib + 'static>(&mut self, item: M) -> &mut Self {
if item.is_prelude() {
self.prelude.push(Rc::new(item));
} else {
self.commands.push(Rc::new(item));
}
self
}
pub fn concat(&mut self, other: &Self) -> &mut Self {
self.prelude.append(&mut other.prelude.clone());
self.commands.append(&mut other.commands.clone());
self
}
fn collect_data(&self) -> (json::Value, Vec<bool>) {
let mut has_data =
vec![false; self.prelude.len() + self.commands.len()];
let data: Vec<json::Value> =
self.prelude.iter()
.chain(self.commands.iter())
.zip(has_data.iter_mut())
.flat_map(|(item, item_has_data)| {
let maybe_data = item.data();
*item_has_data = maybe_data.is_some();
maybe_data
})
.collect();
(json::Value::Array(data), has_data)
}
fn build_script<P>(&self, datafile: P, has_data: &[bool]) -> String
where P: AsRef<Path>
{
let mut script =
format!("\
import json\n\
datafile = open(\"{}\", \"r\")\n\
alldata = json.loads(datafile.read())\n\
datafile.close()\n",
datafile.as_ref().display(),
);
if self.prelude.is_empty() {
script.push_str(PRELUDE);
script.push_str(INIT);
}
let mut data_count: usize = 0;
let iter =
self.prelude.iter()
.chain(self.commands.iter())
.zip(has_data);
for (item, has_data) in iter {
if *has_data {
script.push_str(
&format!("data = alldata[{}]\n", data_count));
data_count += 1;
}
script.push_str(&item.py_cmd());
script.push('\n');
}
script
}
pub fn code(&self, mode: Run) -> String {
let mut tmp_json = get_temp_fname();
tmp_json.set_extension("json");
let (_, has_data) = self.collect_data();
let mut script = self.build_script(&tmp_json, &has_data);
match mode {
Run::Show => {
script.push_str("\nplt.show()");
},
Run::Save(outfile) => {
script.push_str(
&format!("\nfig.savefig(\"{}\")", outfile.display()));
}
Run::SaveShow(outfile) => {
script.push_str(
&format!("\nfig.savefig(\"{}\")", outfile.display()));
script.push_str("\nplt.show()");
},
Run::Debug => { },
Run::Build => { },
}
script
}
pub fn run(&self, mode: Run) -> MplResult<()> {
let tmp = get_temp_fname();
let mut tmp_json = tmp.clone();
tmp_json.set_extension("json");
let mut tmp_py = tmp.clone();
tmp_py.set_extension("py");
let (data, has_data) = self.collect_data();
let mut script = self.build_script(&tmp_json, &has_data);
match mode {
Run::Show => {
script.push_str("\nplt.show()");
},
Run::Save(outfile) => {
script.push_str(
&format!("\nfig.savefig(\"{}\")", outfile.display()));
},
Run::SaveShow(outfile) => {
script.push_str(
&format!("\nfig.savefig(\"{}\")", outfile.display()));
script.push_str("\nplt.show()");
},
Run::Debug => { },
Run::Build => { return Ok(()); },
}
let mut data_file = TempFile::new(&tmp_json)?;
data_file.write_all(json::to_string(&data)?.as_bytes())?;
data_file.flush()?;
let mut script_file = TempFile::new(&tmp_py)?;
script_file.write_all(script.as_bytes())?;
script_file.flush()?;
let res =
process::Command::new("python3")
.arg(format!("{}", tmp_py.display()))
.output()?;
if res.status.success() {
Ok(())
} else {
let stdout: String =
res.stdout.into_iter().map(char::from).collect();
let stderr: String =
res.stderr.into_iter().map(char::from).collect();
Err(MplError::PyError(stdout, stderr))
}
}
pub fn show(&self) -> MplResult<()> { self.run(Run::Show) }
pub fn save<P: AsRef<Path>>(&self, path: P) -> MplResult<()> {
self.run(Run::Save(path.as_ref().to_path_buf()))
}
pub fn saveshow<P: AsRef<Path>>(&self, path: P) -> MplResult<()> {
self.run(Run::SaveShow(path.as_ref().to_path_buf()))
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GSPos {
pub i: Range<usize>,
pub j: Range<usize>,
pub sharex: Option<usize>,
pub sharey: Option<usize>,
}
impl GSPos {
pub fn new(i: Range<usize>, j: Range<usize>) -> Self {
Self { i, j, sharex: None, sharey: None }
}
pub fn new_shared(
i: Range<usize>,
j: Range<usize>,
sharex: Option<usize>,
sharey: Option<usize>,
) -> Self
{
Self { i, j, sharex, sharey }
}
pub fn share(mut self, axis: Axis2, target: Option<usize>) -> Self {
match axis {
Axis2::X => { self.sharex = target; },
Axis2::Y => { self.sharey = target; },
Axis2::Both => { self.sharex = target; self.sharey = target; },
}
self
}
pub fn sharex(mut self, target: Option<usize>) -> Self {
self.sharex = target;
self
}
pub fn sharey(mut self, target: Option<usize>) -> Self {
self.sharey = target;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Run {
Show,
Save(PathBuf),
SaveShow(PathBuf),
Debug,
Build,
}
impl<T: Matplotlib + 'static> From<T> for Mpl {
fn from(item: T) -> Self {
let mut mpl = Self::default();
mpl.then(item);
mpl
}
}
impl std::ops::BitAnd<Mpl> for Mpl {
type Output = Mpl;
fn bitand(mut self, mut rhs: Mpl) -> Self::Output {
self.prelude.append(&mut rhs.prelude);
self.commands.append(&mut rhs.commands);
self
}
}
impl std::ops::BitAndAssign<Mpl> for Mpl {
fn bitand_assign(&mut self, mut rhs: Mpl) {
self.prelude.append(&mut rhs.prelude);
self.commands.append(&mut rhs.commands);
}
}
impl<T> std::ops::BitAnd<T> for Mpl
where T: Matplotlib + 'static
{
type Output = Mpl;
fn bitand(mut self, rhs: T) -> Self::Output {
self.then(rhs);
self
}
}
impl<T> std::ops::BitAndAssign<T> for Mpl
where T: Matplotlib + 'static
{
fn bitand_assign(&mut self, rhs: T) {
self.then(rhs);
}
}
impl std::ops::BitAnd<Run> for Mpl {
type Output = Mpl;
fn bitand(self, mode: Run) -> Self::Output {
match self.run(mode) {
Ok(_) => self,
Err(err) => { panic!("error in Mpl::bitand: {err}"); },
}
}
}
impl std::ops::BitAndAssign<Run> for Mpl {
fn bitand_assign(&mut self, mode: Run) {
match self.run(mode) {
Ok(_) => { },
Err(err) => { panic!("error in Mpl::bitand_assign: {err}"); },
}
}
}
impl std::ops::BitOr<Run> for Mpl {
type Output = ();
fn bitor(self, mode: Run) -> Self::Output {
match self.run(mode) {
Ok(_) => (),
Err(err) => { panic!("error in Mpl::bitor: {err}"); },
}
}
}
impl std::ops::BitOr<Run> for &Mpl {
type Output = ();
fn bitor(self, mode: Run) -> Self::Output {
match self.run(mode) {
Ok(_) => (),
Err(err) => { panic!("error in Mpl::bitor: {err}"); },
}
}
}