Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

improve trillium client: prettify json, add file saving #14

Merged
merged 1 commit into from
Dec 24, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 36 additions & 34 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 4 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ async-channel = "2.1.1"
async-fs = "2.1.0"
async-global-executor = "2.4.1"
async-io = "2.2.2"
atty = "0.2.14"
blocking = "1.5.1"
broadcaster = "1.0.0"
env_logger = "0.10.1"
Expand All @@ -32,13 +31,16 @@ trillium-router = "0.3.5"
trillium-static = { version = "0.4.0", features = ["smol"] }
trillium-websockets = { version = "0.5.2", features = ["json"] }
trillium-smol = "0.3.1"
trillium-client = "0.4.9"
trillium-client = { version = "0.4.9", features = ["json"] }
trillium-logger = "0.4.3"
trillium-html-rewriter = "0.3.0"
url = "2.5.0"
clap = { version = "4.4.11", features = ["derive", "env"] }
clap-verbosity-flag = "2.1.1"
colored = "2.1.0"
colored_json = "4.1.0"
mime = "0.3.17"
size = "0.4.1"

[target.'cfg(unix)'.dependencies]
signal-hook = "0.3.17"
Expand Down
133 changes: 103 additions & 30 deletions src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use trillium_native_tls::NativeTlsConfig;
use trillium_rustls::RustlsConfig;
use trillium_smol::ClientConfig;
use url::{self, Url};

#[derive(Parser, Debug)]
pub struct ClientCli {
#[arg(value_parser = parse_method_case_insensitive)]
Expand All @@ -34,6 +35,10 @@ pub struct ClientCli {
#[arg(short, long, verbatim_doc_comment)]
file: Option<PathBuf>,

/// write the body to a file
#[arg(short, long, verbatim_doc_comment)]
output_file: Option<Option<PathBuf>>,

/// provide a request body on the command line
///
/// example:
Expand Down Expand Up @@ -71,24 +76,79 @@ impl ClientCli {
conn.request_headers().extend(self.headers.clone());

if let Some(path) = &self.file {
let metadata = async_fs::metadata(path)
.await
.unwrap_or_else(|e| panic!("could not read file {:?} ({})", path, e));

let file = async_fs::File::open(path)
.await
.unwrap_or_else(|e| panic!("could not read file {:?} ({})", path, e));

let metadata = file.metadata().await.unwrap();

conn.set_request_body(Body::new_streaming(file, Some(metadata.len())))
} else if let Some(body) = &self.body {
conn.set_request_body(body.clone())
} else if atty::isnt(atty::Stream::Stdin) {
} else if !std::io::stdin().is_terminal() {
conn.set_request_body(Body::new_streaming(Unblock::new(std::io::stdin()), None))
}

conn
}

async fn output_body(&self, conn: &mut Conn) {
if let Some(file) = &self.output_file {
let filename = file
.clone()
.unwrap_or_else(|| conn.url().path_segments().unwrap().last().unwrap().into());
if filename.to_str().map_or(true, |f| f.is_empty()) {
eprintln!("specify a filename for this url");
std::process::exit(-1);
}
let bytes_written = futures_lite::io::copy(
&mut conn.response_body(),
async_fs::File::create(&filename).await.unwrap(),
)
.await
.unwrap();
if matches!(self.verbose.log_level(), Some(level) if level >= Level::Warn) {
println!(
"Wrote {} to {}",
bytes(bytes_written).italic().bright_blue(),
filename.to_string_lossy().italic()
);
}
} else {
let mime: Option<mime::Mime> = conn
.response_headers()
.get_str(trillium_client::KnownHeaderName::ContentType)
.and_then(|ct| ct.parse().ok());
let suffix_or_subtype =
mime.map(|m| m.suffix().unwrap_or_else(|| m.subtype()).to_string());

match suffix_or_subtype.as_deref() {
Some("json") => {
let body = conn.response_json::<serde_json::Value>().await.unwrap();
if std::io::stdout().is_terminal() {
println!("{}", colored_json::to_colored_json_auto(&body).unwrap());
} else {
println!("{}", serde_json::to_string_pretty(&body).unwrap());
}
}

_ => {
match futures_lite::io::copy(
&mut conn.response_body(),
&mut Unblock::new(std::io::stdout()),
)
.await
{
Err(e) if e.kind() == std::io::ErrorKind::WriteZero => {}
other => {
other.unwrap();
}
}
}
}
}
}

pub fn run(self) {
futures_lite::future::block_on(async move {
env_logger::Builder::new()
Expand All @@ -111,20 +171,38 @@ impl ClientCli {

if std::io::stdout().is_terminal() {
let status = conn.status().unwrap_or(Status::NotFound);
println!(
"{}: {}",
"Status".italic(),
if status.is_client_error() {
status.to_string().yellow()
} else if status.is_server_error() {
status.to_string().bright_red()
} else {
status.to_string().bright_green()
}
);
if status != Status::Ok
|| matches!(self.verbose.log_level(), Some(level) if level >= Level::Warn)
{
println!(
"{}: {}",
"Status".italic().bright_blue(),
if status.is_client_error() {
status.to_string().yellow()
} else if status.is_server_error() {
status.to_string().bright_red()
} else {
status.to_string().green()
}
.bold()
);
}

match self.verbose.log_level() {
Some(level) if level >= Level::Warn => {
println!("{}: {}", "Url".italic().bright_blue(), conn.url().as_str());
println!(
"{}: {}",
"Method".italic().bright_blue(),
conn.method().as_str().bold()
);
// if let Some(peer_addr) = conn.peer_addr() {
// println!(
// "{}: {}",
// "Peer Address".italic().bright_blue(),
// peer_addr.to_string()
// );
// }
println!("\n{}", "Request Headers".bold().underline());
print_headers(conn.request_headers());
println!("\n{}", "Response Headers".bold().underline());
Expand All @@ -134,21 +212,9 @@ impl ClientCli {
}
_ => {}
}

futures_lite::io::copy(
&mut conn.response_body(),
&mut Unblock::new(std::io::stdout()),
)
.await
.unwrap();
} else {
futures_lite::io::copy(
&mut conn.response_body(),
&mut Unblock::new(std::io::stdout()),
)
.await
.unwrap();
}

self.output_body(&mut conn).await
});
}
}
Expand Down Expand Up @@ -211,3 +277,10 @@ fn parse_header(s: &str) -> Result<(String, String), String> {
.ok_or_else(|| format!("invalid KEY=value: no `=` found in `{}`", s))?;
Ok((String::from(&s[..pos]), String::from(&s[pos + 1..])))
}

fn bytes(bytes: u64) -> String {
size::Size::from_bytes(bytes)
.format()
.with_base(size::Base::Base10)
.to_string()
}
Loading