enable streaming in random binary/chars (#15361)

# User-Facing Changes

- `random binary` and `random chars` now stream, reducing memory usage
  and allowing interruption with ctrl-c
This commit is contained in:
Solomon 2025-03-20 18:51:22 +00:00 committed by GitHub
parent dd56c813f9
commit 3fe355c4a6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 63 additions and 23 deletions

View File

@ -1,5 +1,5 @@
use super::byte_stream::{random_byte_stream, RandomDistribution};
use nu_engine::command_prelude::*; use nu_engine::command_prelude::*;
use rand::{thread_rng, RngCore};
#[derive(Clone)] #[derive(Clone)]
pub struct RandomBinary; pub struct RandomBinary;
@ -57,12 +57,12 @@ impl Command for RandomBinary {
}), }),
}?; }?;
let mut rng = thread_rng(); Ok(random_byte_stream(
RandomDistribution::Binary,
let mut out = vec![0u8; length]; length,
rng.fill_bytes(&mut out); call.head,
engine_state.signals().clone(),
Ok(Value::binary(out, call.head).into_pipeline_data()) ))
} }
fn examples(&self) -> Vec<Example> { fn examples(&self) -> Vec<Example> {

View File

@ -0,0 +1,49 @@
use nu_engine::command_prelude::*;
use nu_protocol::Signals;
use rand::{
distributions::{Alphanumeric, Standard},
thread_rng, Rng,
};
pub(super) enum RandomDistribution {
Binary,
Alphanumeric,
}
pub(super) fn random_byte_stream(
distribution: RandomDistribution,
length: usize,
span: Span,
signals: Signals,
) -> PipelineData {
let stream_type = match distribution {
RandomDistribution::Binary => ByteStreamType::Binary,
RandomDistribution::Alphanumeric => ByteStreamType::String,
};
const OUTPUT_CHUNK_SIZE: usize = 8192;
let mut remaining_bytes = length;
PipelineData::ByteStream(
ByteStream::from_fn(span, signals.clone(), stream_type, move |out| {
if remaining_bytes == 0 || signals.interrupted() {
return Ok(false);
}
let bytes_to_write = std::cmp::min(remaining_bytes, OUTPUT_CHUNK_SIZE);
let rng = thread_rng();
let byte_iter: Box<dyn Iterator<Item = u8>> = match distribution {
RandomDistribution::Binary => Box::new(rng.sample_iter(Standard)),
RandomDistribution::Alphanumeric => Box::new(rng.sample_iter(Alphanumeric)),
};
out.extend(byte_iter.take(bytes_to_write));
remaining_bytes -= bytes_to_write;
Ok(true)
})
.with_known_size(Some(length as u64)),
None,
)
.with_span(span)
}

View File

@ -1,8 +1,5 @@
use super::byte_stream::{random_byte_stream, RandomDistribution};
use nu_engine::command_prelude::*; use nu_engine::command_prelude::*;
use rand::{
distributions::{Alphanumeric, Distribution},
thread_rng,
};
const DEFAULT_CHARS_LENGTH: usize = 25; const DEFAULT_CHARS_LENGTH: usize = 25;
@ -71,7 +68,6 @@ fn chars(
stack: &mut Stack, stack: &mut Stack,
call: &Call, call: &Call,
) -> Result<PipelineData, ShellError> { ) -> Result<PipelineData, ShellError> {
let span = call.head;
let length: Option<Value> = call.get_flag(engine_state, stack, "length")?; let length: Option<Value> = call.get_flag(engine_state, stack, "length")?;
let length = if let Some(length_val) = length { let length = if let Some(length_val) = length {
match length_val { match length_val {
@ -97,17 +93,11 @@ fn chars(
DEFAULT_CHARS_LENGTH DEFAULT_CHARS_LENGTH
}; };
let mut rng = thread_rng(); Ok(random_byte_stream(
RandomDistribution::Alphanumeric,
let random_string = Alphanumeric length,
.sample_iter(&mut rng) call.head,
.take(length) engine_state.signals().clone(),
.map(char::from)
.collect::<String>();
Ok(PipelineData::Value(
Value::string(random_string, span),
None,
)) ))
} }

View File

@ -1,5 +1,6 @@
mod binary; mod binary;
mod bool; mod bool;
mod byte_stream;
mod chars; mod chars;
mod dice; mod dice;
mod float; mod float;