rune/modules/
capture_io.rs

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::{Address, Memory, Output, VmError};
27use crate::{ContextError, Module};
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            write!(o.inner.lock(), "{m}").map_err(VmError::panic)
39        })
40        .build()?;
41
42    let o = io.clone();
43
44    module
45        .function("println", move |m: &str| {
46            writeln!(o.inner.lock(), "{m}").map_err(VmError::panic)
47        })
48        .build()?;
49
50    let o = io.clone();
51
52    module
53        .raw_function("dbg", move |stack, addr, args, output| {
54            let mut o = o.inner.lock();
55            dbg_impl(&mut o, stack, addr, args, output)
56        })
57        .build()?;
58
59    Ok(module)
60}
61
62/// Type which captures output from rune scripts.
63#[derive(Default, Clone)]
64pub struct CaptureIo {
65    inner: Arc<Mutex<Vec<u8>>>,
66}
67
68impl CaptureIo {
69    /// Construct a new capture.
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    /// Test if capture is empty.
75    pub fn is_empty(&self) -> bool {
76        self.inner.lock().is_empty()
77    }
78
79    /// Drain all captured I/O that has been written to output functions.
80    pub fn drain(&self) -> Vec<u8> {
81        let mut o = self.inner.lock();
82        take(&mut *o)
83    }
84
85    cfg_std! {
86        /// Drain all captured I/O that has been written to output functions into
87        /// the given [Write].
88        ///
89        /// [Write]: std::io::Write
90        pub fn drain_into<O>(&self, mut out: O) -> std::io::Result<()>
91        where
92            O: std::io::Write,
93        {
94            let mut o = self.inner.lock();
95            out.write_all(o.as_slice())?;
96            o.clear();
97            Ok(())
98        }
99    }
100
101    /// Drain all captured I/O that has been written to output functions and try
102    /// to decode as UTF-8.
103    pub fn drain_utf8(&self) -> Result<String, FromUtf8Error> {
104        String::from_utf8(self.drain())
105    }
106}
107
108fn dbg_impl(
109    o: &mut Vec<u8>,
110    memory: &mut dyn Memory,
111    addr: Address,
112    args: usize,
113    out: Output,
114) -> Result<(), VmError> {
115    for value in memory.slice_at(addr, args)? {
116        writeln!(o, "{value:?}").map_err(VmError::panic)?;
117    }
118
119    memory.store(out, ())?;
120    Ok(())
121}
rune/modules/capture_io.rs

rune/modules/
capture_io.rs
