pycall/
lib.rs

1use std::fmt::{Display, Error, Formatter};
2use std::io::Write;
3
4pub trait AsPythonLitteral {
5    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result;
6}
7
8macro_rules! as_py_lit_impl {
9    ($t: ty, $fmt_str: expr) => {
10        impl AsPythonLitteral for $t {
11            fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
12                write!(f, $fmt_str, &self)
13            }
14        }
15    };
16}
17
18as_py_lit_impl!(str, "\"\"\"{}\"\"\"");
19as_py_lit_impl!(String, "\"\"\"{}\"\"\"");
20as_py_lit_impl!(u8, "{}");
21as_py_lit_impl!(u16, "{}");
22as_py_lit_impl!(u32, "{}");
23as_py_lit_impl!(u64, "{}");
24as_py_lit_impl!(u128, "{}");
25as_py_lit_impl!(usize, "{}");
26as_py_lit_impl!(i8, "{}");
27as_py_lit_impl!(i16, "{}");
28as_py_lit_impl!(i32, "{}");
29as_py_lit_impl!(i64, "{}");
30as_py_lit_impl!(i128, "{}");
31as_py_lit_impl!(isize, "{}");
32
33impl AsPythonLitteral for f32 {
34    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
35        if self.is_nan() {
36            write!(f, "float('nan')")
37        } else {
38            write!(f, "{:.6e}", &self)
39        }
40    }
41}
42
43impl AsPythonLitteral for f64 {
44    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
45        if self.is_nan() {
46            write!(f, "float('nan')")
47        } else {
48            write!(f, "{:.6e}", &self)
49        }
50    }
51}
52
53impl<T: AsPythonLitteral> AsPythonLitteral for [T] {
54    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
55        write!(f, "[")?;
56        for x in self.iter() {
57            write!(f, "{},", PythonLiteral(x))?;
58        }
59        write!(f, "]")
60    }
61}
62
63impl<T: AsPythonLitteral> AsPythonLitteral for Vec<T> {
64    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
65        write!(f, "[")?;
66        for x in self.iter() {
67            write!(f, "{},", PythonLiteral(x))?;
68        }
69        write!(f, "]")
70    }
71}
72
73impl<K: AsPythonLitteral, V: AsPythonLitteral> AsPythonLitteral
74    for std::collections::HashMap<K, V>
75{
76    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
77        write!(f, "{{")?;
78        for (k, v) in self.iter() {
79            write!(f, "{}:{},", PythonLiteral(k), PythonLiteral(v))?;
80        }
81        write!(f, "}}")
82    }
83}
84
85#[derive(Copy, Clone, Debug)]
86pub struct Indents(pub isize);
87
88impl std::fmt::Display for Indents {
89    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
90        for _ in 0..self.0 {
91            write!(f, "\t")?;
92        }
93        Ok(())
94    }
95}
96
97struct PythonLiteral<'l, T: AsPythonLitteral + ?Sized>(pub &'l T);
98impl<'l, T: AsPythonLitteral + ?Sized> Display for PythonLiteral<'l, T> {
99    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
100        self.0.fmt(f)
101    }
102}
103
104pub struct JoinGuard<T>(Option<std::thread::JoinHandle<T>>);
105
106impl<T> JoinGuard<T> {
107    pub fn new() -> Self {
108        JoinGuard(None)
109    }
110
111    pub fn spawn<F: FnOnce() -> T>(f: F) -> Self
112    where
113        T: Send + 'static,
114        F: Send + 'static,
115    {
116        JoinGuard(Some(std::thread::spawn(f)))
117    }
118
119    pub fn join(mut self) -> Result<T, Box<dyn std::any::Any + Send>>
120    where
121        T: std::any::Any + Send + 'static,
122    {
123        self.0.take().unwrap().join()
124    }
125
126    pub fn detach(mut self) -> Option<std::thread::JoinHandle<T>> {
127        self.0.take()
128    }
129}
130
131impl<T> Drop for JoinGuard<T> {
132    fn drop(&mut self) {
133        if let Some(handle) = self.0.take() {
134            handle.join();
135        }
136    }
137}
138
139/// An instance of code generation unit.
140/// It really is just a file with dedicated APIs to write Python into it.
141/// Most importantly: it manages indentation for you.
142pub struct PythonProgram {
143    file: tempfile::NamedTempFile,
144    indents: Indents,
145}
146impl PythonProgram {
147    /// Creates a named temp file to store the generated python program
148    pub fn new() -> PythonProgram {
149        PythonProgram {
150            file: tempfile::NamedTempFile::new().unwrap(),
151            indents: Indents(0),
152        }
153    }
154
155    pub fn save_as<P: AsRef<std::path::Path>>(&self, path: P) -> Result<u64, std::io::Error> {
156        std::fs::copy(self.file.path(), path)
157    }
158
159    /// Runs the program using python3
160    pub fn run(&self) -> Result<std::process::Output, std::io::Error> {
161        std::process::Command::new("python3")
162            .arg(self.file.path())
163            .output()
164    }
165
166    /// Spawns a thread to run the program using python3.
167    /// The returned JoinGuard ensures that the program will be ran to completion.
168    pub fn background_run(self) -> JoinGuard<Result<std::process::Output, std::io::Error>> {
169        JoinGuard::spawn(move || self.run())
170    }
171
172    /// Ensures that the internal file has been flushed. Typically not necessary.
173    pub fn flush(&mut self) -> &mut Self {
174        self.file.flush().unwrap();
175        self
176    }
177
178    /// Moves the indentation level by `n`. However, I recommend using the dedicated functions when possible/
179    pub fn indent(&mut self, n: isize) -> &mut Self {
180        if n >= 0 {
181            self.indents.0 += n as isize
182        } else {
183            self.indents.0 -= n as isize
184        }
185        self
186    }
187
188    /// Removes one indentation level from the cursor.
189    /// You should call this whenever you're done with a scope.
190    pub fn end_block(&mut self) -> &mut Self {
191        self.indent(-1)
192    }
193
194    /// Writes a line assigning `value` formatted as a python literal to `name`
195    pub fn define_variable<T: AsPythonLitteral + ?Sized>(
196        &mut self,
197        name: &str,
198        value: &T,
199    ) -> &mut Self {
200        writeln!(
201            &mut self.file,
202            "{}{} = {}",
203            self.indents,
204            name,
205            PythonLiteral(value)
206        )
207        .unwrap();
208        self
209    }
210
211    /// Writes an import statement for your `dependency`
212    pub fn import(&mut self, dependency: &str) -> &mut Self {
213        writeln!(&mut self.file, "{}import {}", self.indents, dependency).unwrap();
214        self
215    }
216
217    /// Writes an import statement for your `dependency` as `rename`
218    pub fn import_as(&mut self, dependency: &str, rename: &str) -> &mut Self {
219        writeln!(
220            &mut self.file,
221            "{}import {} as {}",
222            self.indents, dependency, rename
223        )
224        .unwrap();
225        self
226    }
227
228    /// Writes whatever line you passed it, indented at the proper level.
229    pub fn write_line(&mut self, line: &str) -> &mut Self {
230        writeln!(&mut self.file, "{}{}", self.indents, line).unwrap();
231        self
232    }
233
234    /// Writes an if, using your condition as a test, and increments indentation.
235    pub fn r#if(&mut self, condition: &str) -> &mut Self {
236        writeln!(&mut self.file, "{}if {}:", self.indents, condition).unwrap();
237        self.indent(1)
238    }
239    /// Decrements indentation, writes an elif, using your condition as a test, and increments indentation.
240    pub fn elif(&mut self, condition: &str) -> &mut Self {
241        self.indent(-1);
242        writeln!(&mut self.file, "{}elif {}:", self.indents, condition).unwrap();
243        self.indent(1)
244    }
245    /// Decrements indentation, writes an else, using your condition as a test, and increments indentation.
246    pub fn r#else(&mut self) -> &mut Self {
247        self.indent(-1).write_line("else:").indent(1)
248    }
249
250    /// Writes "for `range`:", and increments indentation.
251    pub fn r#for(&mut self, range: &str) -> &mut Self {
252        writeln!(&mut self.file, "{}for {}:", self.indents, range).unwrap();
253        self.indent(1)
254    }
255
256    /// Writes a while, using your condition as a test, and increments indentation.
257    pub fn r#while(&mut self, condition: &str) -> &mut Self {
258        writeln!(&mut self.file, "{}while {}:", self.indents, condition).unwrap();
259        self.indent(1)
260    }
261}
262
263impl Write for PythonProgram {
264    fn write(&mut self, buf: &[u8]) -> Result<usize, std::io::Error> {
265        self.file.write(buf)
266    }
267
268    fn flush(&mut self) -> Result<(), std::io::Error> {
269        self.file.flush()
270    }
271}
272
273impl Display for PythonProgram {
274    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
275        use std::io::BufRead;
276        let read_file = std::fs::File::open(self.file.path()).unwrap();
277        let reader = std::io::BufReader::new(read_file);
278        for line in reader.lines() {
279            writeln!(f, "{}", line.unwrap())?
280        }
281        Ok(())
282    }
283}
284
285pub trait MatPlotLib {
286    fn import_pyplot_as_plt(&mut self) -> &mut Self;
287    fn plot_y<Y: AsPythonLitteral>(&mut self, y: &Y) -> &mut Self;
288    fn plot_xy<X: AsPythonLitteral, Y: AsPythonLitteral>(&mut self, x: &X, y: &Y) -> &mut Self;
289    fn plot_xyargs<X: AsPythonLitteral, Y: AsPythonLitteral>(
290        &mut self,
291        x: &X,
292        y: &Y,
293        args: &str,
294    ) -> &mut Self;
295    fn semilogy_y<Y: AsPythonLitteral>(&mut self, y: &Y) -> &mut Self;
296    fn semilogy_xy<X: AsPythonLitteral, Y: AsPythonLitteral>(&mut self, x: &X, y: &Y) -> &mut Self;
297    fn semilogy_xyargs<X: AsPythonLitteral, Y: AsPythonLitteral>(
298        &mut self,
299        x: &X,
300        y: &Y,
301        args: &str,
302    ) -> &mut Self;
303    fn show(&mut self) -> &mut Self;
304}
305
306impl MatPlotLib for PythonProgram {
307    fn import_pyplot_as_plt(&mut self) -> &mut Self {
308        self.import_as("matplotlib.pyplot", "plt")
309    }
310
311    fn plot_y<Y: AsPythonLitteral>(&mut self, y: &Y) -> &mut Self {
312        self.write_line(&format!("plt.plot({})", PythonLiteral(y)))
313    }
314
315    fn plot_xy<X: AsPythonLitteral, Y: AsPythonLitteral>(&mut self, x: &X, y: &Y) -> &mut Self {
316        self.write_line(&format!(
317            "plt.plot({},{})",
318            PythonLiteral(x),
319            PythonLiteral(y)
320        ))
321    }
322
323    fn plot_xyargs<X: AsPythonLitteral, Y: AsPythonLitteral>(
324        &mut self,
325        x: &X,
326        y: &Y,
327        args: &str,
328    ) -> &mut Self {
329        self.write_line(&format!(
330            "plt.plot({},{},{})",
331            PythonLiteral(x),
332            PythonLiteral(y),
333            args
334        ))
335    }
336
337    fn semilogy_y<Y: AsPythonLitteral>(&mut self, y: &Y) -> &mut Self {
338        self.write_line(&format!("plt.semilogy({})", PythonLiteral(y)))
339    }
340
341    fn semilogy_xy<X: AsPythonLitteral, Y: AsPythonLitteral>(&mut self, x: &X, y: &Y) -> &mut Self {
342        self.write_line(&format!(
343            "plt.semilogy({},{})",
344            PythonLiteral(x),
345            PythonLiteral(y)
346        ))
347    }
348
349    fn semilogy_xyargs<X: AsPythonLitteral, Y: AsPythonLitteral>(
350        &mut self,
351        x: &X,
352        y: &Y,
353        args: &str,
354    ) -> &mut Self {
355        self.write_line(&format!(
356            "plt.semilogy({},{},{})",
357            PythonLiteral(x),
358            PythonLiteral(y),
359            args
360        ))
361    }
362
363    fn show(&mut self) -> &mut Self {
364        self.write_line("plt.show()")
365    }
366}
367
368pub mod plots {
369    use crate::{AsPythonLitteral, PythonLiteral, PythonProgram};
370    use std::io::Write;
371
372    pub fn plot_xyargs<X: AsPythonLitteral, Y: AsPythonLitteral>(
373        x: &X,
374        y: &Y,
375        args: &str,
376    ) -> Result<std::process::Output, std::io::Error> {
377        let mut program = PythonProgram::new();
378        program.import_as("matplotlib.pyplot", "plt");
379        writeln!(
380            &program.file,
381            "plt.plot({}, {}, {})",
382            PythonLiteral(x),
383            PythonLiteral(y),
384            PythonLiteral(args)
385        );
386        program.write_line("plt.show()").run()
387    }
388
389    pub fn plot_xy<X: AsPythonLitteral, Y: AsPythonLitteral>(
390        x: &X,
391        y: &Y,
392    ) -> Result<std::process::Output, std::io::Error> {
393        let mut program = PythonProgram::new();
394        program.import_as("matplotlib.pyplot", "plt");
395        writeln!(
396            &program.file,
397            "plt.plot({}, {})",
398            PythonLiteral(x),
399            PythonLiteral(y),
400        );
401        program.write_line("plt.show()").run()
402    }
403
404    pub fn plot_y<Y: AsPythonLitteral>(y: &Y) -> Result<std::process::Output, std::io::Error> {
405        let mut program = PythonProgram::new();
406        program.import_as("matplotlib.pyplot", "plt");
407        writeln!(&program.file, "plt.plot({})", PythonLiteral(y));
408        program.write_line("plt.show()").run()
409    }
410}
411
412#[macro_export]
413macro_rules! plot {
414    ($y: expr) => {
415        pycall::plots::plot_y($y)
416    };
417    ($x: expr, $y: expr) => {
418        pycall::plots::plot_xy($x, $y)
419    };
420    ($x: expr, $y: expr, $args: expr) => {
421        pycall::plots::plot_xyargs($x, $y, $args)
422    };
423}
424
425#[test]
426fn run() {
427    let join = std::thread::spawn(|| quick_plot(&(-50..50).map(|x| (-x * x)).collect::<Vec<_>>()));
428    let mut program = PythonProgram::new();
429    program
430        .write_line("import matplotlib.pyplot as plt")
431        .define_variable(
432            "hello",
433            &(-50..50).map(|x| (x * x) as f64).collect::<Vec<_>>(),
434        )
435        .write_line("print(hello)")
436        .write_line("plt.plot(hello)")
437        .write_line("plt.show()");
438    println!("program: {}\r\n{}", program.file.path().display(), &program);
439    let output = program.run().unwrap();
440    join.join();
441}
pycall/lib.rs

pycall/
lib.rs
