mirror of https://github.com/sharkdp/bat
Add `$LESSOPEN` and `$LESSCLOSE` support (#2444)
parent
3abc0c0fc4
commit
e32ad0b048
@ -0,0 +1,390 @@
|
||||
#![cfg(feature = "lessopen")]
|
||||
|
||||
use std::convert::TryFrom;
|
||||
use std::env;
|
||||
use std::fs::File;
|
||||
use std::io::{BufRead, BufReader, Cursor, Read, Write};
|
||||
use std::path::PathBuf;
|
||||
use std::str;
|
||||
|
||||
use clircle::{Clircle, Identifier};
|
||||
use os_str_bytes::RawOsString;
|
||||
use run_script::{IoOptions, ScriptOptions};
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::input::{Input, InputKind, InputReader, OpenedInput, OpenedInputKind};
|
||||
|
||||
/// Preprocess files and/or stdin using $LESSOPEN and $LESSCLOSE
|
||||
pub(crate) struct LessOpenPreprocessor {
|
||||
lessopen: String,
|
||||
lessclose: Option<String>,
|
||||
command_options: ScriptOptions,
|
||||
kind: LessOpenKind,
|
||||
/// Whether or not data piped via stdin is to be preprocessed
|
||||
preprocess_stdin: bool,
|
||||
}
|
||||
|
||||
enum LessOpenKind {
|
||||
Piped,
|
||||
PipedIgnoreExitCode,
|
||||
TempFile,
|
||||
}
|
||||
|
||||
impl LessOpenPreprocessor {
|
||||
/// Create a new instance of LessOpenPreprocessor
|
||||
/// Will return Ok if and only if $LESSOPEN is set
|
||||
pub(crate) fn new() -> Result<LessOpenPreprocessor> {
|
||||
let lessopen = env::var("LESSOPEN")?;
|
||||
|
||||
// "||" means pipe directly to bat without making a temporary file
|
||||
// Also, if preprocessor output is empty and exit code is zero, use the empty output
|
||||
// Otherwise, if output is empty and exit code is nonzero, use original file contents
|
||||
let (kind, lessopen) = if lessopen.starts_with("||") {
|
||||
(LessOpenKind::Piped, lessopen.chars().skip(2).collect())
|
||||
// "|" means pipe, but ignore exit code, always using preprocessor output
|
||||
} else if lessopen.starts_with('|') {
|
||||
(
|
||||
LessOpenKind::PipedIgnoreExitCode,
|
||||
lessopen.chars().skip(1).collect(),
|
||||
)
|
||||
// If neither appear, write output to a temporary file and read from that
|
||||
} else {
|
||||
(LessOpenKind::TempFile, lessopen)
|
||||
};
|
||||
|
||||
// "-" means that stdin is preprocessed along with files and may appear alongside "|" and "||"
|
||||
let (stdin, lessopen) = if lessopen.starts_with('-') {
|
||||
(true, lessopen.chars().skip(1).collect())
|
||||
} else {
|
||||
(false, lessopen)
|
||||
};
|
||||
|
||||
let mut command_options = ScriptOptions::new();
|
||||
command_options.runner = env::var("SHELL").ok();
|
||||
command_options.input_redirection = IoOptions::Pipe;
|
||||
|
||||
Ok(Self {
|
||||
lessopen: lessopen.replacen("%s", "$1", 1),
|
||||
lessclose: env::var("LESSCLOSE")
|
||||
.ok()
|
||||
.map(|str| str.replacen("%s", "$1", 1).replacen("%s", "$2", 1)),
|
||||
command_options,
|
||||
kind,
|
||||
preprocess_stdin: stdin,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn open<'a, R: BufRead + 'a>(
|
||||
&self,
|
||||
input: Input<'a>,
|
||||
mut stdin: R,
|
||||
stdout_identifier: Option<&Identifier>,
|
||||
) -> Result<OpenedInput<'a>> {
|
||||
let (lessopen_stdout, path_str, kind) = match input.kind {
|
||||
InputKind::OrdinaryFile(ref path) => {
|
||||
let path_str = match path.to_str() {
|
||||
Some(str) => str,
|
||||
None => return input.open(stdin, stdout_identifier),
|
||||
};
|
||||
|
||||
let (exit_code, lessopen_stdout, _) = match run_script::run(
|
||||
&self.lessopen,
|
||||
&vec![path_str.to_string()],
|
||||
&self.command_options,
|
||||
) {
|
||||
Ok(output) => output,
|
||||
Err(_) => return input.open(stdin, stdout_identifier),
|
||||
};
|
||||
|
||||
if self.fall_back_to_original_file(&lessopen_stdout, exit_code) {
|
||||
return input.open(stdin, stdout_identifier);
|
||||
}
|
||||
|
||||
(
|
||||
RawOsString::from_string(lessopen_stdout),
|
||||
path_str.to_string(),
|
||||
OpenedInputKind::OrdinaryFile(path.to_path_buf()),
|
||||
)
|
||||
}
|
||||
InputKind::StdIn => {
|
||||
if self.preprocess_stdin {
|
||||
if let Some(stdout) = stdout_identifier {
|
||||
let input_identifier = Identifier::try_from(clircle::Stdio::Stdin)
|
||||
.map_err(|e| format!("Stdin: Error identifying file: {}", e))?;
|
||||
if stdout.surely_conflicts_with(&input_identifier) {
|
||||
return Err("IO circle detected. The input from stdin is also an output. Aborting to avoid infinite loop.".into());
|
||||
}
|
||||
}
|
||||
|
||||
// stdin isn't Clone, so copy it to a cloneable buffer
|
||||
let mut stdin_buffer = Vec::new();
|
||||
stdin.read_to_end(&mut stdin_buffer).unwrap();
|
||||
|
||||
let mut lessopen_handle = match run_script::spawn(
|
||||
&self.lessopen,
|
||||
&vec!["-".to_string()],
|
||||
&self.command_options,
|
||||
) {
|
||||
Ok(handle) => handle,
|
||||
Err(_) => {
|
||||
return input.open(stdin, stdout_identifier);
|
||||
}
|
||||
};
|
||||
|
||||
if lessopen_handle
|
||||
.stdin
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.write_all(&stdin_buffer.clone())
|
||||
.is_err()
|
||||
{
|
||||
return input.open(stdin, stdout_identifier);
|
||||
}
|
||||
|
||||
let lessopen_output = match lessopen_handle.wait_with_output() {
|
||||
Ok(output) => output,
|
||||
Err(_) => {
|
||||
return input.open(Cursor::new(stdin_buffer), stdout_identifier);
|
||||
}
|
||||
};
|
||||
|
||||
if lessopen_output.stdout.is_empty()
|
||||
&& (!lessopen_output.status.success()
|
||||
|| matches!(self.kind, LessOpenKind::PipedIgnoreExitCode))
|
||||
{
|
||||
return input.open(Cursor::new(stdin_buffer), stdout_identifier);
|
||||
}
|
||||
|
||||
(
|
||||
RawOsString::assert_from_raw_vec(lessopen_output.stdout),
|
||||
"-".to_string(),
|
||||
OpenedInputKind::StdIn,
|
||||
)
|
||||
} else {
|
||||
return input.open(stdin, stdout_identifier);
|
||||
}
|
||||
}
|
||||
InputKind::CustomReader(_) => {
|
||||
return input.open(stdin, stdout_identifier);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(OpenedInput {
|
||||
kind,
|
||||
reader: InputReader::new(BufReader::new(
|
||||
if matches!(self.kind, LessOpenKind::TempFile) {
|
||||
// Remove newline at end of temporary file path returned by $LESSOPEN
|
||||
let stdout = match lessopen_stdout.strip_suffix("\n") {
|
||||
Some(stripped) => stripped.to_owned(),
|
||||
None => lessopen_stdout,
|
||||
};
|
||||
|
||||
let stdout = stdout.into_os_string();
|
||||
|
||||
let file = match File::open(PathBuf::from(&stdout)) {
|
||||
Ok(file) => file,
|
||||
Err(_) => {
|
||||
return input.open(stdin, stdout_identifier);
|
||||
}
|
||||
};
|
||||
|
||||
Preprocessed {
|
||||
kind: PreprocessedKind::TempFile(file),
|
||||
lessclose: self.lessclose.clone(),
|
||||
command_args: vec![path_str, stdout.to_str().unwrap().to_string()],
|
||||
command_options: self.command_options.clone(),
|
||||
}
|
||||
} else {
|
||||
Preprocessed {
|
||||
kind: PreprocessedKind::Piped(Cursor::new(lessopen_stdout.into_raw_vec())),
|
||||
lessclose: self.lessclose.clone(),
|
||||
command_args: vec![path_str, "-".to_string()],
|
||||
command_options: self.command_options.clone(),
|
||||
}
|
||||
},
|
||||
)),
|
||||
metadata: input.metadata,
|
||||
description: input.description,
|
||||
})
|
||||
}
|
||||
|
||||
fn fall_back_to_original_file(&self, lessopen_output: &str, exit_code: i32) -> bool {
|
||||
lessopen_output.is_empty()
|
||||
&& (exit_code != 0 || matches!(self.kind, LessOpenKind::PipedIgnoreExitCode))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
/// For testing purposes only
|
||||
/// Create an instance of LessOpenPreprocessor with specified valued for $LESSOPEN and $LESSCLOSE
|
||||
fn mock_new(lessopen: Option<&str>, lessclose: Option<&str>) -> Result<LessOpenPreprocessor> {
|
||||
if let Some(command) = lessopen {
|
||||
env::set_var("LESSOPEN", command)
|
||||
} else {
|
||||
env::remove_var("LESSOPEN")
|
||||
}
|
||||
|
||||
if let Some(command) = lessclose {
|
||||
env::set_var("LESSCLOSE", command)
|
||||
} else {
|
||||
env::remove_var("LESSCLOSE")
|
||||
}
|
||||
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
enum PreprocessedKind {
|
||||
Piped(Cursor<Vec<u8>>),
|
||||
TempFile(File),
|
||||
}
|
||||
|
||||
impl Read for PreprocessedKind {
|
||||
fn read(&mut self, buf: &mut [u8]) -> std::result::Result<usize, std::io::Error> {
|
||||
match self {
|
||||
PreprocessedKind::Piped(data) => data.read(buf),
|
||||
PreprocessedKind::TempFile(data) => data.read(buf),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Preprocessed {
|
||||
kind: PreprocessedKind,
|
||||
lessclose: Option<String>,
|
||||
command_args: Vec<String>,
|
||||
command_options: ScriptOptions,
|
||||
}
|
||||
|
||||
impl Read for Preprocessed {
|
||||
fn read(&mut self, buf: &mut [u8]) -> std::result::Result<usize, std::io::Error> {
|
||||
self.kind.read(buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Preprocessed {
|
||||
fn drop(&mut self) {
|
||||
if let Some(ref command) = self.lessclose {
|
||||
self.command_options.output_redirection = IoOptions::Inherit;
|
||||
|
||||
run_script::run(command, &self.command_args, &self.command_options)
|
||||
.expect("failed to run $LESSCLOSE to clean up file");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
// All tests here are serial because they all involve reading and writing environment variables
|
||||
// Running them in parallel causes these tests and some others to randomly fail
|
||||
use serial_test::serial;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Reset environment variables after each test as a precaution
|
||||
fn reset_env_vars() {
|
||||
env::remove_var("LESSOPEN");
|
||||
env::remove_var("LESSCLOSE");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_just_lessopen() -> Result<()> {
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("|batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(preprocessor.lessclose.is_none());
|
||||
|
||||
reset_env_vars();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_just_lessclose() -> Result<()> {
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(None, Some("lessclose.sh %s %s"));
|
||||
|
||||
assert!(preprocessor.is_err());
|
||||
|
||||
reset_env_vars();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_both_lessopen_and_lessclose() -> Result<()> {
|
||||
let preprocessor =
|
||||
LessOpenPreprocessor::mock_new(Some("lessopen.sh %s"), Some("lessclose.sh %s %s"))?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "lessopen.sh $1");
|
||||
assert_eq!(preprocessor.lessclose.unwrap(), "lessclose.sh $1 $2");
|
||||
|
||||
reset_env_vars();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn test_lessopen_prefixes() -> Result<()> {
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(matches!(preprocessor.kind, LessOpenKind::TempFile));
|
||||
assert!(!preprocessor.preprocess_stdin);
|
||||
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("|batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(matches!(
|
||||
preprocessor.kind,
|
||||
LessOpenKind::PipedIgnoreExitCode
|
||||
));
|
||||
assert!(!preprocessor.preprocess_stdin);
|
||||
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("||batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(matches!(preprocessor.kind, LessOpenKind::Piped));
|
||||
assert!(!preprocessor.preprocess_stdin);
|
||||
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("-batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(matches!(preprocessor.kind, LessOpenKind::TempFile));
|
||||
assert!(preprocessor.preprocess_stdin);
|
||||
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("|-batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(matches!(
|
||||
preprocessor.kind,
|
||||
LessOpenKind::PipedIgnoreExitCode
|
||||
));
|
||||
assert!(preprocessor.preprocess_stdin);
|
||||
|
||||
let preprocessor = LessOpenPreprocessor::mock_new(Some("||-batpipe %s"), None)?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "batpipe $1");
|
||||
assert!(matches!(preprocessor.kind, LessOpenKind::Piped));
|
||||
assert!(preprocessor.preprocess_stdin);
|
||||
|
||||
reset_env_vars();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn replace_part_of_argument() -> Result<()> {
|
||||
let preprocessor =
|
||||
LessOpenPreprocessor::mock_new(Some("|echo File:%s"), Some("echo File:%s Temp:%s"))?;
|
||||
|
||||
assert_eq!(preprocessor.lessopen, "echo File:$1");
|
||||
assert_eq!(preprocessor.lessclose.unwrap(), "echo File:$1 Temp:$2");
|
||||
|
||||
reset_env_vars();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue