1use std::{
2 collections::HashMap,
3 fs,
4 io::{ self, Write },
5 ops::Range,
6 path::{ Path, PathBuf },
7 process,
8 rc::Rc,
9};
10use rand::distributions::{ Alphanumeric, DistString };
11use serde_json as json;
12use thiserror::Error;
13use crate::commands::Axis2;
14
15#[derive(Debug, Error)]
16pub enum MplError {
17 #[error("IO error: {0}")]
18 IOError(#[from] io::Error),
19
20 #[error("serialization error: {0}")]
21 JsonError(#[from] json::Error),
22
23 #[error("script error:\nstdout:\n{0}\nstderr:\n{1}")]
24 PyError(String, String),
25}
26pub type MplResult<T> = Result<T, MplError>;
27
28pub const PRELUDE: &str
51= "\
52import datetime
53import io
54import json
55import os
56import random
57import sys
58import matplotlib
59matplotlib.use(\"QtAgg\")
60import matplotlib.path as mpath
61import matplotlib.patches as mpatches
62import matplotlib.pyplot as plt
63import matplotlib.cm as mcm
64import matplotlib.colors as mcolors
65import matplotlib.collections as mcollections
66import matplotlib.ticker as mticker
67import matplotlib.image as mimage
68from mpl_toolkits.mplot3d import axes3d
69import numpy as np
70";
71
72pub const INIT: &str
79= "\
80fig, ax = plt.subplots()
81";
82
83pub trait Matplotlib: std::fmt::Debug {
85 fn is_prelude(&self) -> bool;
88
89 fn data(&self) -> Option<json::Value>;
92
93 fn py_cmd(&self) -> String;
101}
102
103pub trait AsPy {
105 fn as_py(&self) -> String;
106}
107
108impl AsPy for bool {
109 fn as_py(&self) -> String { if *self { "True" } else { "False" }.into() }
110}
111
112impl AsPy for i32 {
113 fn as_py(&self) -> String { self.to_string() }
114}
115
116impl AsPy for f64 {
117 fn as_py(&self) -> String { self.to_string() }
118}
119
120impl AsPy for String {
121 fn as_py(&self) -> String { format!("\"{self}\"") }
122}
123
124impl AsPy for &str {
125 fn as_py(&self) -> String { format!("\"{self}\"") }
126}
127
128#[derive(Clone, Debug, PartialEq)]
130pub enum PyValue {
131 Bool(bool),
133 Int(i32),
135 Float(f64),
137 Str(String),
139 List(Vec<PyValue>),
141 Dict(HashMap<String, PyValue>),
143 Var(String),
147 None
149}
150
151impl From<bool> for PyValue {
152 fn from(b: bool) -> Self { Self::Bool(b) }
153}
154
155impl From<i32> for PyValue {
156 fn from(i: i32) -> Self { Self::Int(i) }
157}
158
159impl From<f64> for PyValue {
160 fn from(f: f64) -> Self { Self::Float(f) }
161}
162
163impl From<String> for PyValue {
164 fn from(s: String) -> Self { Self::Str(s) }
165}
166
167impl From<&str> for PyValue {
168 fn from(s: &str) -> Self { Self::Str(s.into()) }
169}
170
171impl<T> From<&T> for PyValue
172where T: Clone + Into<PyValue>
173{
174 fn from(x: &T) -> Self { x.clone().into() }
175}
176
177impl From<Vec<PyValue>> for PyValue {
178 fn from(l: Vec<PyValue>) -> Self { Self::List(l) }
179}
180
181impl From<HashMap<String, PyValue>> for PyValue {
182 fn from(d: HashMap<String, PyValue>) -> Self { Self::Dict(d) }
183}
184
185impl<T: Into<PyValue>> FromIterator<T> for PyValue {
186 fn from_iter<I>(iter: I) -> Self
187 where I: IntoIterator<Item = T>
188 {
189 Self::list(iter)
190 }
191}
192
193impl<S: Into<String>, T: Into<PyValue>> FromIterator<(S, T)> for PyValue {
194 fn from_iter<I>(iter: I) -> Self
195 where I: IntoIterator<Item = (S, T)>
196 {
197 Self::dict(iter)
198 }
199}
200
201impl PyValue {
202 pub fn list<I, T>(items: I) -> Self
204 where
205 I: IntoIterator<Item = T>,
206 T: Into<PyValue>,
207 {
208 Self::List(items.into_iter().map(|item| item.into()).collect())
209 }
210
211 pub fn dict<I, S, T>(items: I) -> Self
213 where
214 I: IntoIterator<Item = (S, T)>,
215 S: Into<String>,
216 T: Into<PyValue>,
217 {
218 Self::Dict(
219 items.into_iter().map(|(s, v)| (s.into(), v.into())).collect())
220 }
221}
222
223impl AsPy for PyValue {
224 fn as_py(&self) -> String {
225 match self {
226 Self::Bool(b) => if *b { "True".into() } else { "False".into() },
227 Self::Int(i) => format!("{i}"),
228 Self::Float(f) => format!("{f}"),
229 Self::Str(s) => format!("\"{s}\""),
230 Self::List(l) => {
231 let n = l.len();
232 let mut out = String::from("[");
233 for (k, v) in l.iter().enumerate() {
234 out.push_str(&v.as_py());
235 if k < n - 1 { out.push_str(", "); }
236 }
237 out.push(']');
238 out
239 },
240 Self::Dict(d) => {
241 let n = d.len();
242 let mut out = String::from("{");
243 for (j, (k, v)) in d.iter().enumerate() {
244 out.push_str(&format!("\"{}\": {}", k, v.as_py()));
245 if j < n - 1 { out.push_str(", "); }
246 }
247 out.push('}');
248 out
249 },
250 Self::Var(v) => v.clone(),
251 Self::None => "None".into(),
252 }
253 }
254}
255
256#[derive(Clone, Debug, PartialEq)]
258pub struct Opt(pub String, pub PyValue);
259
260impl<T: Into<PyValue>> From<(&str, T)> for Opt {
261 fn from(kv: (&str, T)) -> Self { Self(kv.0.into(), kv.1.into()) }
262}
263
264impl<T: Into<PyValue>> From<(String, T)> for Opt {
265 fn from(kv: (String, T)) -> Self { Self(kv.0, kv.1.into()) }
266}
267
268impl Opt {
269 pub fn new<T>(key: &str, val: T) -> Self
271 where T: Into<PyValue>
272 {
273 Self(key.into(), val.into())
274 }
275}
276
277pub fn opt<T>(key: &str, val: T) -> Opt
279where T: Into<PyValue>
280{
281 Opt::new(key, val)
282}
283
284impl AsPy for Opt {
285 fn as_py(&self) -> String { format!("{}={}", self.0, self.1.as_py()) }
286}
287
288impl AsPy for Vec<Opt> {
289 fn as_py(&self) -> String {
290 let n = self.len();
291 let mut out = String::new();
292 for (k, opt) in self.iter().enumerate() {
293 out.push_str(&opt.as_py());
294 if k < n - 1 { out.push_str(", "); }
295 }
296 out
297 }
298}
299
300pub trait MatplotlibOpts: Matplotlib {
302 fn kwarg<T: Into<PyValue>>(&mut self, key: &str, val: T) -> &mut Self;
304
305 fn o<T: Into<PyValue>>(mut self, key: &str, val: T) -> Self
307 where Self: Sized
308 {
309 self.kwarg(key, val);
310 self
311 }
312
313 fn oo<I>(mut self, opts: I) -> Self
315 where
316 I: IntoIterator<Item = Opt>,
317 Self: Sized,
318 {
319 opts.into_iter().for_each(|Opt(key, val)| { self.kwarg(&key, val); });
320 self
321 }
322}
323
324fn get_temp_fname() -> PathBuf {
325 std::env::temp_dir()
326 .join(Alphanumeric.sample_string(&mut rand::thread_rng(), 15))
327}
328
329#[derive(Debug)]
330struct TempFile(PathBuf, Option<fs::File>);
331
332impl TempFile {
333 fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
334 let path = path.as_ref().to_path_buf();
335 fs::OpenOptions::new()
336 .create(true)
337 .append(false)
338 .truncate(true)
339 .write(true)
340 .open(&path)
341 .map(|file| Self(path, Some(file)))
342 }
343}
344
345impl Write for TempFile {
346 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
347 if let Some(file) = self.1.as_mut() {
348 file.write(buf)
349 } else {
350 Ok(0)
351 }
352 }
353
354 fn flush(&mut self) -> io::Result<()> {
355 if let Some(file) = self.1.as_mut() {
356 file.flush()
357 } else {
358 Ok(())
359 }
360 }
361}
362
363impl Drop for TempFile {
364 fn drop(&mut self) {
365 drop(self.1.take());
366 fs::remove_file(&self.0).ok();
367 }
368}
369
370#[derive(Clone, Debug, Default)]
401pub struct Mpl {
402 prelude: Vec<Rc<dyn Matplotlib + 'static>>,
403 commands: Vec<Rc<dyn Matplotlib + 'static>>,
404}
405impl Mpl {
408 pub fn new() -> Self { Self::default() }
416
417 pub fn new_3d<I>(opts: I) -> Self
425 where I: IntoIterator<Item = Opt>
426 {
427 let opts: Vec<Opt> = opts.into_iter().collect();
428 Self::default()
429 & crate::commands::DefPrelude
430 & crate::commands::Init3D { opts: opts.into_iter().collect() }
431 }
432
433 pub fn new_3d_with<I, F>(opts: I, f: F) -> Self
436 where
437 I: IntoIterator<Item = Opt>,
438 F: FnOnce(Mpl) -> Mpl,
439 {
440 f(Self::default() & crate::commands::DefPrelude)
441 & crate::commands::Init3D { opts: opts.into_iter().collect() }
442 }
443
444 pub fn new_grid<I>(nrows: usize, ncols: usize, opts: I) -> Self
454 where I: IntoIterator<Item = Opt>
455 {
456 let opts: Vec<Opt> = opts.into_iter().collect();
457 Self::default()
458 & crate::commands::DefPrelude
459 & crate::commands::InitGrid {
460 nrows,
461 ncols,
462 opts: opts.into_iter().collect(),
463 }
464 }
465
466 pub fn new_grid_with<I, F>(nrows: usize, ncols: usize, opts: I, f: F)
469 -> Self
470 where
471 I: IntoIterator<Item = Opt>,
472 F: FnOnce(Mpl) -> Mpl,
473 {
474 f(Self::default() & crate::commands::DefPrelude)
475 & crate::commands::InitGrid {
476 nrows,
477 ncols,
478 opts: opts.into_iter().collect(),
479 }
480 }
481
482 pub fn new_gridspec<I, P>(gridspec_kw: I, positions: P) -> Self
493 where
494 I: IntoIterator<Item = Opt>,
495 P: IntoIterator<Item = GSPos>,
496 {
497 Self::default()
498 & crate::commands::DefPrelude
499 & crate::commands::init_gridspec(gridspec_kw, positions)
500 }
501
502 pub fn new_gridspec_with<I, P, F>(gridspec_kw: I, positions: P, f: F)
505 -> Self
506 where
507 I: IntoIterator<Item = Opt>,
508 P: IntoIterator<Item = GSPos>,
509 F: FnOnce(Mpl) -> Mpl,
510 {
511 f(Self::default() & crate::commands::DefPrelude)
512 & crate::commands::init_gridspec(gridspec_kw, positions)
513 }
514
515 pub fn then<M: Matplotlib + 'static>(&mut self, item: M) -> &mut Self {
517 if item.is_prelude() {
518 self.prelude.push(Rc::new(item));
519 } else {
520 self.commands.push(Rc::new(item));
521 }
522 self
523 }
524
525 pub fn concat(&mut self, other: &Self) -> &mut Self {
530 self.prelude.append(&mut other.prelude.clone());
531 self.commands.append(&mut other.commands.clone());
532 self
533 }
534
535 fn collect_data(&self) -> (json::Value, Vec<bool>) {
536 let mut has_data =
537 vec![false; self.prelude.len() + self.commands.len()];
538 let data: Vec<json::Value> =
539 self.prelude.iter()
540 .chain(self.commands.iter())
541 .zip(has_data.iter_mut())
542 .flat_map(|(item, item_has_data)| {
543 let maybe_data = item.data();
544 *item_has_data = maybe_data.is_some();
545 maybe_data
546 })
547 .collect();
548 (json::Value::Array(data), has_data)
549 }
550
551 fn build_script<P>(&self, datafile: P, has_data: &[bool]) -> String
552 where P: AsRef<Path>
553 {
554 let mut script =
555 format!("\
556 import json\n\
557 datafile = open(\"{}\", \"r\")\n\
558 alldata = json.loads(datafile.read())\n\
559 datafile.close()\n",
560 datafile.as_ref().display(),
561 );
562 if self.prelude.is_empty() {
563 script.push_str(PRELUDE);
564 script.push_str(INIT);
565 }
566 let mut data_count: usize = 0;
567 let iter =
568 self.prelude.iter()
569 .chain(self.commands.iter())
570 .zip(has_data);
571 for (item, has_data) in iter {
572 if *has_data {
573 script.push_str(
574 &format!("data = alldata[{}]\n", data_count));
575 data_count += 1;
576 }
577 script.push_str(&item.py_cmd());
578 script.push('\n');
579 }
580 script
581 }
582
583 pub fn code(&self, mode: Run) -> String {
585 let mut tmp_json = get_temp_fname();
586 tmp_json.set_extension("json");
587 let (_, has_data) = self.collect_data();
588 let mut script = self.build_script(&tmp_json, &has_data);
589 match mode {
590 Run::Show => {
591 script.push_str("\nplt.show()");
592 },
593 Run::Save(outfile) => {
594 script.push_str(
595 &format!("\nfig.savefig(\"{}\")", outfile.display()));
596 }
597 Run::SaveShow(outfile) => {
598 script.push_str(
599 &format!("\nfig.savefig(\"{}\")", outfile.display()));
600 script.push_str("\nplt.show()");
601 },
602 Run::Debug => { },
603 Run::Build => { },
604 }
605 script
606 }
607
608 pub fn run(&self, mode: Run) -> MplResult<()> {
610 let tmp = get_temp_fname();
611 let mut tmp_json = tmp.clone();
612 tmp_json.set_extension("json");
613 let mut tmp_py = tmp.clone();
614 tmp_py.set_extension("py");
615 let (data, has_data) = self.collect_data();
616 let mut script = self.build_script(&tmp_json, &has_data);
617 match mode {
618 Run::Show => {
619 script.push_str("\nplt.show()");
620 },
621 Run::Save(outfile) => {
622 script.push_str(
623 &format!("\nfig.savefig(\"{}\")", outfile.display()));
624 },
625 Run::SaveShow(outfile) => {
626 script.push_str(
627 &format!("\nfig.savefig(\"{}\")", outfile.display()));
628 script.push_str("\nplt.show()");
629 },
630 Run::Debug => { },
631 Run::Build => { return Ok(()); },
632 }
633 let mut data_file = TempFile::new(&tmp_json)?;
634 data_file.write_all(json::to_string(&data)?.as_bytes())?;
635 data_file.flush()?;
636 let mut script_file = TempFile::new(&tmp_py)?;
637 script_file.write_all(script.as_bytes())?;
638 script_file.flush()?;
639 let res =
640 process::Command::new("python3")
641 .arg(format!("{}", tmp_py.display()))
642 .output()?;
643 if res.status.success() {
644 Ok(())
645 } else {
646 let stdout: String =
647 res.stdout.into_iter().map(char::from).collect();
648 let stderr: String =
649 res.stderr.into_iter().map(char::from).collect();
650 Err(MplError::PyError(stdout, stderr))
651 }
652 }
653
654 pub fn show(&self) -> MplResult<()> { self.run(Run::Show) }
656
657 pub fn save<P: AsRef<Path>>(&self, path: P) -> MplResult<()> {
659 self.run(Run::Save(path.as_ref().to_path_buf()))
660 }
661
662 pub fn saveshow<P: AsRef<Path>>(&self, path: P) -> MplResult<()> {
664 self.run(Run::SaveShow(path.as_ref().to_path_buf()))
665 }
666}
667
668#[derive(Clone, Debug, PartialEq, Eq)]
686pub struct GSPos {
687 pub i: Range<usize>,
689 pub j: Range<usize>,
691 pub sharex: Option<usize>,
693 pub sharey: Option<usize>,
695}
696
697impl GSPos {
698 pub fn new(i: Range<usize>, j: Range<usize>) -> Self {
700 Self { i, j, sharex: None, sharey: None }
701 }
702
703 pub fn new_shared(
705 i: Range<usize>,
706 j: Range<usize>,
707 sharex: Option<usize>,
708 sharey: Option<usize>,
709 ) -> Self
710 {
711 Self { i, j, sharex, sharey }
712 }
713
714 pub fn share(mut self, axis: Axis2, target: Option<usize>) -> Self {
716 match axis {
717 Axis2::X => { self.sharex = target; },
718 Axis2::Y => { self.sharey = target; },
719 Axis2::Both => { self.sharex = target; self.sharey = target; },
720 }
721 self
722 }
723
724 pub fn sharex(mut self, target: Option<usize>) -> Self {
726 self.sharex = target;
727 self
728 }
729
730 pub fn sharey(mut self, target: Option<usize>) -> Self {
732 self.sharey = target;
733 self
734 }
735}
736
737#[derive(Clone, Debug, PartialEq, Eq)]
740pub enum Run {
741 Show,
743 Save(PathBuf),
745 SaveShow(PathBuf),
747 Debug,
750 Build,
752}
753
754impl<T: Matplotlib + 'static> From<T> for Mpl {
755 fn from(item: T) -> Self {
756 let mut mpl = Self::default();
757 mpl.then(item);
758 mpl
759 }
760}
761
762impl std::ops::BitAnd<Mpl> for Mpl {
763 type Output = Mpl;
764
765 fn bitand(mut self, mut rhs: Mpl) -> Self::Output {
766 self.prelude.append(&mut rhs.prelude);
767 self.commands.append(&mut rhs.commands);
768 self
769 }
770}
771
772impl std::ops::BitAndAssign<Mpl> for Mpl {
773 fn bitand_assign(&mut self, mut rhs: Mpl) {
774 self.prelude.append(&mut rhs.prelude);
775 self.commands.append(&mut rhs.commands);
776 }
777}
778
779impl<T> std::ops::BitAnd<T> for Mpl
780where T: Matplotlib + 'static
781{
782 type Output = Mpl;
783
784 fn bitand(mut self, rhs: T) -> Self::Output {
785 self.then(rhs);
786 self
787 }
788}
789
790impl<T> std::ops::BitAndAssign<T> for Mpl
791where T: Matplotlib + 'static
792{
793 fn bitand_assign(&mut self, rhs: T) {
794 self.then(rhs);
795 }
796}
797
798impl std::ops::BitAnd<Run> for Mpl {
799 type Output = Mpl;
800
801 fn bitand(self, mode: Run) -> Self::Output {
802 match self.run(mode) {
803 Ok(_) => self,
804 Err(err) => { panic!("error in Mpl::bitand: {err}"); },
805 }
806 }
807}
808
809impl std::ops::BitAndAssign<Run> for Mpl {
810 fn bitand_assign(&mut self, mode: Run) {
811 match self.run(mode) {
812 Ok(_) => { },
813 Err(err) => { panic!("error in Mpl::bitand_assign: {err}"); },
814 }
815 }
816}
817
818impl std::ops::BitOr<Run> for Mpl {
819 type Output = ();
820
821 fn bitor(self, mode: Run) -> Self::Output {
822 match self.run(mode) {
823 Ok(_) => (),
824 Err(err) => { panic!("error in Mpl::bitor: {err}"); },
825 }
826 }
827}
828
829impl std::ops::BitOr<Run> for &Mpl {
830 type Output = ();
831
832 fn bitor(self, mode: Run) -> Self::Output {
833 match self.run(mode) {
834 Ok(_) => (),
835 Err(err) => { panic!("error in Mpl::bitor: {err}"); },
836 }
837 }
838}
839