nushell/crates/nu-protocol/src/value/stream.rs

148 lines
3.9 KiB
Rust
Raw Normal View History

2021-09-08 04:26:57 +02:00
use crate::*;
2021-10-28 06:13:10 +02:00
use std::{
fmt::Debug,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
};
2021-09-08 04:26:57 +02:00
/// A single buffer of binary data streamed over multiple parts. Optionally contains ctrl-c that can be used
/// to break the stream.
pub struct ByteStream {
pub stream: Box<dyn Iterator<Item = Vec<u8>> + Send + 'static>,
pub ctrlc: Option<Arc<AtomicBool>>,
}
impl ByteStream {
pub fn into_vec(self) -> Vec<u8> {
self.flatten().collect::<Vec<u8>>()
}
}
impl Debug for ByteStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ByteStream").finish()
}
}
impl Iterator for ByteStream {
type Item = Vec<u8>;
fn next(&mut self) -> Option<Self::Item> {
if let Some(ctrlc) = &self.ctrlc {
if ctrlc.load(Ordering::SeqCst) {
None
} else {
self.stream.next()
}
} else {
self.stream.next()
}
}
}
/// A single string streamed over multiple parts. Optionally contains ctrl-c that can be used
/// to break the stream.
pub struct StringStream {
pub stream: Box<dyn Iterator<Item = Result<String, ShellError>> + Send + 'static>,
pub ctrlc: Option<Arc<AtomicBool>>,
}
impl StringStream {
pub fn into_string(self, separator: &str) -> Result<String, ShellError> {
let mut output = String::new();
let mut first = true;
for s in self.stream {
output.push_str(&s?);
if !first {
output.push_str(separator);
} else {
first = false;
}
}
Ok(output)
}
pub fn from_stream(
input: impl Iterator<Item = Result<String, ShellError>> + Send + 'static,
ctrlc: Option<Arc<AtomicBool>>,
) -> StringStream {
StringStream {
stream: Box::new(input),
ctrlc,
}
}
}
impl Debug for StringStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StringStream").finish()
}
}
impl Iterator for StringStream {
type Item = Result<String, ShellError>;
fn next(&mut self) -> Option<Self::Item> {
if let Some(ctrlc) = &self.ctrlc {
if ctrlc.load(Ordering::SeqCst) {
None
} else {
self.stream.next()
}
} else {
self.stream.next()
}
}
}
2021-11-03 01:26:09 +01:00
/// A potentially infinite stream of values, optinally with a mean to send a Ctrl-C signal to stop
/// the stream from continuing.
///
/// In practice, a "stream" here means anything which can be iterated and produce Values as it iterates.
/// Like other iterators in Rust, observing values from this stream will drain the items as you view them
/// and the stream cannot be replayed.
2021-10-28 06:13:10 +02:00
pub struct ValueStream {
pub stream: Box<dyn Iterator<Item = Value> + Send + 'static>,
pub ctrlc: Option<Arc<AtomicBool>>,
}
2021-09-08 04:26:57 +02:00
impl ValueStream {
pub fn into_string(self, separator: &str, config: &Config) -> String {
self.map(|x: Value| x.into_string(", ", config))
2021-11-09 21:17:37 +01:00
.collect::<Vec<String>>()
.join(separator)
2021-09-08 04:26:57 +02:00
}
2021-10-28 06:13:10 +02:00
pub fn from_stream(
input: impl Iterator<Item = Value> + Send + 'static,
ctrlc: Option<Arc<AtomicBool>>,
) -> ValueStream {
ValueStream {
stream: Box::new(input),
ctrlc,
}
2021-09-08 04:26:57 +02:00
}
}
impl Debug for ValueStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ValueStream").finish()
}
}
impl Iterator for ValueStream {
type Item = Value;
fn next(&mut self) -> Option<Self::Item> {
2021-10-28 06:13:10 +02:00
if let Some(ctrlc) = &self.ctrlc {
if ctrlc.load(Ordering::SeqCst) {
None
} else {
self.stream.next()
}
} else {
self.stream.next()
2021-10-11 20:12:47 +02:00
}
2021-10-01 07:11:49 +02:00
}
}