rune/modules/
capture_io.rs

Help
1//! I/O module capable of capturing what's been written to a buffer.
2//!
3//! # Examples
4//!
5//! ```no_run
6//! use rune::Context;
7//! use rune::modules::capture_io::{self, CaptureIo};
8//!
9//! let io = CaptureIo::new();
10//!
11//! let mut context = rune::Context::with_config(false)?;
12//! context.install(capture_io::module(&io)?)?;
13//! # Ok::<_, rune::ContextError>(())
14//! ```
15
16use core::mem::take;
17
18use ::rust_alloc::sync::Arc;
19
20use parking_lot::Mutex;
21
22use crate as rune;
23use crate::alloc::fmt::TryWrite;
24use crate::alloc::string::FromUtf8Error;
25use crate::alloc::{String, Vec};
26use crate::runtime::{InstAddress, Memory, Output, VmError, VmResult};
27use crate::{ContextError, Module, Value};
28
29/// I/O module capable of capturing what's been written to a buffer.
30#[rune::module(::std::io)]
31pub fn module(io: &CaptureIo) -> Result<Module, ContextError> {
32    let mut module = Module::from_meta(self::module_meta)?;
33
34    let o = io.clone();
35
36    module
37        .function("print", move |m: &str| {
38            match write!(o.inner.lock(), "{}", m) {
39                Ok(()) => VmResult::Ok(()),
40                Err(error) => VmResult::panic(error),
41            }
42        })
43        .build()?;
44
45    let o = io.clone();
46
47    module
48        .function("println", move |m: &str| {
49            match writeln!(o.inner.lock(), "{}", m) {
50                Ok(()) => VmResult::Ok(()),
51                Err(error) => VmResult::panic(error),
52            }
53        })
54        .build()?;
55
56    let o = io.clone();
57
58    module
59        .raw_function("dbg", move |stack, addr, args, output| {
60            let mut o = o.inner.lock();
61            dbg_impl(&mut o, stack, addr, args, output)
62        })
63        .build()?;
64
65    Ok(module)
66}
67
68/// Type which captures output from rune scripts.
69#[derive(Default, Clone)]
70pub struct CaptureIo {
71    inner: Arc<Mutex<Vec<u8>>>,
72}
73
74impl CaptureIo {
75    /// Construct a new capture.
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    /// Test if capture is empty.
81    pub fn is_empty(&self) -> bool {
82        self.inner.lock().is_empty()
83    }
84
85    /// Drain all captured I/O that has been written to output functions.
86    pub fn drain(&self) -> Vec<u8> {
87        let mut o = self.inner.lock();
88        take(&mut *o)
89    }
90
91    cfg_std! {
92        /// Drain all captured I/O that has been written to output functions into
93        /// the given [Write].
94        ///
95        /// [Write]: std::io::Write
96        pub fn drain_into<O>(&self, mut out: O) -> std::io::Result<()>
97        where
98            O: std::io::Write,
99        {
100            let mut o = self.inner.lock();
101            out.write_all(o.as_slice())?;
102            o.clear();
103            Ok(())
104        }
105    }
106
107    /// Drain all captured I/O that has been written to output functions and try
108    /// to decode as UTF-8.
109    pub fn drain_utf8(&self) -> Result<String, FromUtf8Error> {
110        String::from_utf8(self.drain())
111    }
112}
113
114fn dbg_impl(
115    o: &mut Vec<u8>,
116    stack: &mut dyn Memory,
117    addr: InstAddress,
118    args: usize,
119    out: Output,
120) -> VmResult<()> {
121    for value in vm_try!(stack.slice_at(addr, args)) {
122        vm_try!(writeln!(o, "{:?}", value).map_err(VmError::panic));
123    }
124
125    vm_try!(out.store(stack, Value::unit));
126    VmResult::Ok(())
127}
rune/modules/capture_io.rs

rune/modules/
capture_io.rs
