Allow captured stderr saving to file (#6793)

* support redirect stderr to file

* fix test

* fix test

* fix test
This commit is contained in:
WindSoilder 2022-10-20 20:56:44 +08:00 committed by GitHub
parent d37e6ba3b5
commit 10aa86272b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 162 additions and 27 deletions

View File

@ -2,8 +2,10 @@ use nu_engine::CallExt;
use nu_protocol::ast::Call;
use nu_protocol::engine::{Command, EngineState, Stack};
use nu_protocol::{
Category, Example, PipelineData, ShellError, Signature, Spanned, SyntaxShape, Value,
Category, Example, PipelineData, RawStream, ShellError, Signature, Span, Spanned, SyntaxShape,
Value,
};
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
@ -35,6 +37,12 @@ impl Command for Save {
fn signature(&self) -> nu_protocol::Signature {
Signature::build("save")
.required("filename", SyntaxShape::Filepath, "the filename to use")
.named(
"stderr",
SyntaxShape::Filepath,
"the filename used to save stderr, only works with `-r` flag",
Some('e'),
)
.switch("raw", "save file as raw binary", Some('r'))
.switch("append", "append input to the end of the file", Some('a'))
.category(Category::FileSystem)
@ -81,6 +89,35 @@ impl Command for Save {
));
}
};
let stderr_path = call.get_flag::<Spanned<String>>(engine_state, stack, "stderr")?;
let stderr_file = match stderr_path {
None => None,
Some(stderr_path) => {
let stderr_span = stderr_path.span;
let stderr_path = Path::new(&stderr_path.item);
if stderr_path == path {
Some(file.try_clone()?)
} else {
match std::fs::File::create(stderr_path) {
Ok(file) => Some(file),
Err(err) => {
return Ok(PipelineData::Value(
Value::Error {
error: ShellError::GenericError(
"Permission denied".into(),
err.to_string(),
Some(stderr_span),
None,
Vec::new(),
),
},
None,
))
}
}
}
}
};
let ext = if raw {
None
@ -148,33 +185,37 @@ impl Command for Save {
match input {
PipelineData::ExternalStream { stdout: None, .. } => Ok(PipelineData::new(span)),
PipelineData::ExternalStream {
stdout: Some(mut stream),
stdout: Some(stream),
stderr,
..
} => {
let mut writer = BufWriter::new(file);
// delegate a thread to redirect stderr to result.
let handler = stderr.map(|stderr_stream| match stderr_file {
Some(stderr_file) => std::thread::spawn(move || {
stream_to_file(stderr_stream, stderr_file, span)
}),
None => std::thread::spawn(move || {
let _ = stderr_stream.into_bytes();
Ok(PipelineData::new(span))
}),
});
stream
.try_for_each(move |result| {
let buf = match result {
Ok(v) => match v {
Value::String { val, .. } => val.into_bytes(),
Value::Binary { val, .. } => val,
_ => {
return Err(ShellError::UnsupportedInput(
format!("{:?} not supported", v.get_type()),
v.span()?,
));
}
},
Err(err) => return Err(err),
};
if let Err(err) = writer.write(&buf) {
return Err(ShellError::IOError(err.to_string()));
let res = stream_to_file(stream, file, span);
if let Some(h) = handler {
match h.join() {
Err(err) => {
return Err(ShellError::ExternalCommand(
"Fail to receive external commands stderr message".to_string(),
format!("{err:?}"),
span,
))
}
Ok(())
})
.map(|_| PipelineData::new(span))
Ok(res) => res,
}?;
res
} else {
res
}
}
input => match input.into_value(span) {
Value::String { val, .. } => {
@ -237,6 +278,47 @@ impl Command for Save {
example: r#"echo { a: 1, b: 2 } | save foo.json"#,
result: None,
},
Example {
description: "Save a running program's stderr to foo.txt",
example: r#"do -i {} | save foo.txt --stderr foo.txt"#,
result: None,
},
Example {
description: "Save a running program's stderr to separate file",
example: r#"do -i {} | save foo.txt --stderr bar.txt"#,
result: None,
},
]
}
}
fn stream_to_file(
mut stream: RawStream,
file: File,
span: Span,
) -> Result<PipelineData, ShellError> {
let mut writer = BufWriter::new(file);
stream
.try_for_each(move |result| {
let buf = match result {
Ok(v) => match v {
Value::String { val, .. } => val.into_bytes(),
Value::Binary { val, .. } => val,
_ => {
return Err(ShellError::UnsupportedInput(
format!("{:?} not supported", v.get_type()),
v.span()?,
));
}
},
Err(err) => return Err(err),
};
if let Err(err) = writer.write(&buf) {
return Err(ShellError::IOError(err.to_string()));
}
Ok(())
})
.map(|_| PipelineData::new(span))
}

View File

@ -82,3 +82,51 @@ fn save_append_will_not_overwrite_content() {
assert_eq!(actual, "hello world");
})
}
#[test]
fn save_stderr_and_stdout_to_same_file() {
Playground::setup("save_test_5", |dirs, sandbox| {
sandbox.with_files(vec![]);
let expected_file = dirs.test().join("new-file.txt");
nu!(
cwd: dirs.root(),
r#"
let-env FOO = "bar";
let-env BAZ = "ZZZ";
do -i {nu -c 'nu --testbin echo_env FOO; nu --testbin echo_env_stderr BAZ'} | save -r save_test_5/new-file.txt --stderr save_test_5/new-file.txt"#,
);
let actual = file_contents(expected_file);
println!("{}, {}", actual, actual.contains("ZZZ"));
assert!(actual.contains("bar"));
assert!(actual.contains("ZZZ"));
})
}
#[test]
fn save_stderr_and_stdout_to_diff_file() {
Playground::setup("save_test_6", |dirs, sandbox| {
sandbox.with_files(vec![]);
let expected_file = dirs.test().join("log.txt");
let expected_stderr_file = dirs.test().join("err.txt");
nu!(
cwd: dirs.root(),
r#"
let-env FOO = "bar";
let-env BAZ = "ZZZ";
do -i {nu -c 'nu --testbin echo_env FOO; nu --testbin echo_env_stderr BAZ'} | save -r save_test_6/log.txt --stderr save_test_6/err.txt"#,
);
let actual = file_contents(expected_file);
assert!(actual.contains("bar"));
assert!(!actual.contains("ZZZ"));
let actual = file_contents(expected_stderr_file);
assert!(actual.contains("ZZZ"));
assert!(!actual.contains("bar"));
})
}

View File

@ -292,7 +292,8 @@ fn main() -> Result<()> {
if let Some(testbin) = &binary_args.testbin {
// Call out to the correct testbin
match testbin.item.as_str() {
"echo_env" => test_bins::echo_env(),
"echo_env" => test_bins::echo_env(true),
"echo_env_stderr" => test_bins::echo_env(false),
"cococo" => test_bins::cococo(),
"meow" => test_bins::meow(),
"meowb" => test_bins::meowb(),

View File

@ -11,11 +11,15 @@ use nu_protocol::{CliError, PipelineData, Span, Value};
/// Echo's value of env keys from args
/// Example: nu --testbin env_echo FOO BAR
/// If it it's not present echo's nothing
pub fn echo_env() {
pub fn echo_env(to_stdout: bool) {
let args = args();
for arg in args {
if let Ok(v) = std::env::var(arg) {
println!("{}", v);
if to_stdout {
println!("{}", v);
} else {
eprintln!("{}", v);
}
}
}
}