Allow byte streams with unknown type to be compatiable with binary (#12959)

# Description
Currently, this pipeline doesn't work `open --raw file | take 100`,
since the type of the byte stream is `Unknown`, but `take` expects
`Binary` streams. This PR changes commands that expect
`ByteStreamType::Binary` to also work with `ByteStreamType::Unknown`.
This was done by adding two new methods to `ByteStreamType`:
`is_binary_coercible` and `is_string_coercible`. These return true if
the type is `Unknown` or matches the type in the method name.
This commit is contained in:
Ian Manske 2024-05-25 00:54:38 +00:00 committed by GitHub
parent b06f31d3c6
commit c5d716951f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 21 additions and 13 deletions

View File

@ -160,7 +160,7 @@ fn string_helper(
// Just set the type - that should be good enough. There is no guarantee that the data // Just set the type - that should be good enough. There is no guarantee that the data
// within a string stream is actually valid UTF-8. But refuse to do it if it was already set // within a string stream is actually valid UTF-8. But refuse to do it if it was already set
// to binary // to binary
if stream.type_() != ByteStreamType::Binary { if stream.type_().is_string_coercible() {
Ok(PipelineData::ByteStream( Ok(PipelineData::ByteStream(
stream.with_type(ByteStreamType::String), stream.with_type(ByteStreamType::String),
metadata, metadata,

View File

@ -1,4 +1,5 @@
use nu_engine::command_prelude::*; use nu_engine::command_prelude::*;
use std::io::Read;
#[derive(Clone)] #[derive(Clone)]
pub struct First; pub struct First;
@ -171,10 +172,9 @@ fn first_helper(
} }
} }
PipelineData::ByteStream(stream, metadata) => { PipelineData::ByteStream(stream, metadata) => {
if stream.type_() == ByteStreamType::Binary { if stream.type_().is_binary_coercible() {
let span = stream.span(); let span = stream.span();
if let Some(mut reader) = stream.reader() { if let Some(mut reader) = stream.reader() {
use std::io::Read;
if return_single_element { if return_single_element {
// Take a single byte // Take a single byte
let mut byte = [0u8]; let mut byte = [0u8];

View File

@ -1,6 +1,5 @@
use nu_engine::command_prelude::*; use nu_engine::command_prelude::*;
use std::{collections::VecDeque, io::Read};
use std::collections::VecDeque;
#[derive(Clone)] #[derive(Clone)]
pub struct Last; pub struct Last;
@ -161,10 +160,9 @@ impl Command for Last {
} }
} }
PipelineData::ByteStream(stream, ..) => { PipelineData::ByteStream(stream, ..) => {
if stream.type_() == ByteStreamType::Binary { if stream.type_().is_binary_coercible() {
let span = stream.span(); let span = stream.span();
if let Some(mut reader) = stream.reader() { if let Some(mut reader) = stream.reader() {
use std::io::Read;
// Have to be a bit tricky here, but just consume into a VecDeque that we // Have to be a bit tricky here, but just consume into a VecDeque that we
// shrink to fit each time // shrink to fit each time
const TAKE: u64 = 8192; const TAKE: u64 = 8192;

View File

@ -1,4 +1,5 @@
use nu_engine::command_prelude::*; use nu_engine::command_prelude::*;
use std::io::{self, Read};
#[derive(Clone)] #[derive(Clone)]
pub struct Skip; pub struct Skip;
@ -94,12 +95,11 @@ impl Command for Skip {
let input_span = input.span().unwrap_or(call.head); let input_span = input.span().unwrap_or(call.head);
match input { match input {
PipelineData::ByteStream(stream, metadata) => { PipelineData::ByteStream(stream, metadata) => {
if stream.type_() == ByteStreamType::Binary { if stream.type_().is_binary_coercible() {
let span = stream.span(); let span = stream.span();
if let Some(mut reader) = stream.reader() { if let Some(mut reader) = stream.reader() {
use std::io::Read;
// Copy the number of skipped bytes into the sink before proceeding // Copy the number of skipped bytes into the sink before proceeding
std::io::copy(&mut (&mut reader).take(n as u64), &mut std::io::sink()) io::copy(&mut (&mut reader).take(n as u64), &mut io::sink())
.err_span(span)?; .err_span(span)?;
Ok(PipelineData::ByteStream( Ok(PipelineData::ByteStream(
ByteStream::read(reader, call.head, None, ByteStreamType::Binary), ByteStream::read(reader, call.head, None, ByteStreamType::Binary),

View File

@ -1,4 +1,5 @@
use nu_engine::command_prelude::*; use nu_engine::command_prelude::*;
use std::io::Read;
#[derive(Clone)] #[derive(Clone)]
pub struct Take; pub struct Take;
@ -79,9 +80,8 @@ impl Command for Take {
metadata, metadata,
)), )),
PipelineData::ByteStream(stream, metadata) => { PipelineData::ByteStream(stream, metadata) => {
if stream.type_() == ByteStreamType::Binary { if stream.type_().is_binary_coercible() {
if let Some(reader) = stream.reader() { if let Some(reader) = stream.reader() {
use std::io::Read;
// Just take 'rows' bytes off the stream, mimicking the binary behavior // Just take 'rows' bytes off the stream, mimicking the binary behavior
Ok(PipelineData::ByteStream( Ok(PipelineData::ByteStream(
ByteStream::read( ByteStream::read(

View File

@ -114,6 +114,16 @@ impl ByteStreamType {
ByteStreamType::Unknown => "byte stream", ByteStreamType::Unknown => "byte stream",
} }
} }
/// Returns true if the type is `Binary` or `Unknown`
pub fn is_binary_coercible(self) -> bool {
matches!(self, ByteStreamType::Binary | ByteStreamType::Unknown)
}
/// Returns true if the type is `String` or `Unknown`
pub fn is_string_coercible(self) -> bool {
matches!(self, ByteStreamType::String | ByteStreamType::Unknown)
}
} }
impl From<ByteStreamType> for Type { impl From<ByteStreamType> for Type {
@ -483,7 +493,7 @@ impl ByteStream {
/// data would have been valid UTF-8. /// data would have been valid UTF-8.
pub fn into_string(self) -> Result<String, ShellError> { pub fn into_string(self) -> Result<String, ShellError> {
let span = self.span; let span = self.span;
if self.type_ != ByteStreamType::Binary { if self.type_.is_string_coercible() {
let trim = self.stream.is_external(); let trim = self.stream.is_external();
let bytes = self.into_bytes()?; let bytes = self.into_bytes()?;
let mut string = String::from_utf8(bytes).map_err(|err| ShellError::NonUtf8Custom { let mut string = String::from_utf8(bytes).map_err(|err| ShellError::NonUtf8Custom {