buckn/derive_http
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ud, fixed errors

Browse Source
This commit is contained in:
Nathan Buck
2025-07-18 13:34:27 -04:00
parent 0ba23fad0e
commit 41f85afa6c
1 changed files with 58 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

138
src/lib.rs
View File

@ -147,85 +147,79 @@ pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn alpaca_cli(_attr: TokenStream, item: TokenStream) -> TokenStream { pub fn alpaca_cli(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input_enum = parse_macro_input!(item as ItemEnum); let input_enum = parse_macro_input!(item as ItemEnum);
let top_enum_ident = &input_enum.ident; // e.g., "Cmd" let top_enum_ident = &input_enum.ident;
let top_variants = &input_enum.variants; let top_variants = &input_enum.variants;
// For each variant like Alpaca(AlpacaCmd) // Build outer match arms
let outer_match_arms = top_variants.iter().map(|v| { let match_arms: Vec<_> = top_variants.iter().map(|variant| {
let variant_ident = &v.ident; // e.g., "Alpaca" let variant_ident = &variant.ident;
// Extract the inner sub-enum type like AlpacaCmd // Expecting tuple variants like Alpaca(AlpacaCmd)
if let syn::Fields::Unnamed(fields) = &v.fields { let inner_type = match &variant.fields {
if let Some(field) = fields.unnamed.first() { Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
if let syn::Type::Path(inner_type) = &field.ty { match &fields.unnamed.first().unwrap().ty {
let inner_type_ident = &inner_type.path.segments.last().unwrap().ident; syn::Type::Path(p) => p.path.segments.last().unwrap().ident.clone(),
_ => panic!("Expected tuple variant with a type path"),
}
}
_ => panic!("Each variant must be a tuple variant like `Alpaca(AlpacaCmd)`"),
};
// Match arms inside the nested enum (AlpacaCmd) quote! {
let inner_match_arm = quote! { #top_enum_ident::#variant_ident(inner) => {
match #inner_type_ident::parse() { match inner {
#inner_type_ident::Bulk { input } => { #inner_type::Bulk { input } => {
// Bulk: read and parse Vec<inner_type_ident> let mut reader: Box<dyn std::io::Read> = match input {
let mut reader: Box<dyn std::io::Read> = match input { Some(path) => Box::new(std::fs::File::open(path)?),
Some(path) => Box::new(std::fs::File::open(path)?), None => Box::new(std::io::stdin()),
None => Box::new(std::io::stdin()), };
};
let mut buf = String::new();
reader.read_to_string(&mut buf)?;
let queries: Vec<#inner_type_ident> = serde_json::from_str(&buf)?;
use std::sync::Arc; let mut buf = String::new();
let client = Arc::new(awc::Client::default()); reader.read_to_string(&mut buf)?;
let api_keys = Arc::new(crate::load_api_keys()?); let queries: Vec<#inner_type> = serde_json::from_str(&buf)?;
const THREADS: usize = 4; use std::sync::Arc;
let total = queries.len(); let client = Arc::new(awc::Client::default());
let per_thread = total / THREADS; let keys = Arc::new(crate::load_api_keys()?);
let shared_queries = Arc::new(queries);
let mut handles = Vec::new(); const THREADS: usize = 4;
for i in 0..THREADS { let total = queries.len();
let queries = Arc::clone(&shared_queries); let per_thread = std::cmp::max(1, total / THREADS);
let client = Arc::clone(&client); let shared_queries = Arc::new(queries);
let keys = Arc::clone(&api_keys);
let start = i * per_thread;
let end = if i == THREADS - 1 { total } else { start + per_thread };
let handle = std::thread::spawn(move || { let mut handles = Vec::new();
let rt = tokio::runtime::Runtime::new().unwrap(); for i in 0..THREADS {
for q in &queries[start..end] { let queries = Arc::clone(&shared_queries);
rt.block_on(q.send_all(&client, &keys)); let client = Arc::clone(&client);
} let keys = Arc::clone(&keys);
}); let start = i * per_thread;
let end = if i == THREADS - 1 { total } else { start + per_thread };
handles.push(handle); let handle = std::thread::spawn(move || {
let rt = tokio::runtime::Runtime::new().unwrap();
for q in &queries[start..end] {
rt.block_on(q.send_all(&client, &keys)).unwrap();
} }
});
for h in handles { handles.push(handle);
h.join().expect("Thread panicked");
}
}
other => {
let client = awc::Client::default();
let keys = crate::load_api_keys()?;
other.send_all(&client, &keys).await?;
}
} }
};
// Wrap the outer enum match (Cmd::Alpaca(inner)) for h in handles {
return quote! { h.join().expect("Thread panicked");
#top_enum_ident::#variant_ident(inner) => {
#inner_match_arm
} }
}; }
other => {
let client = awc::Client::default();
let keys = crate::load_api_keys()?;
other.send_all(&client, &keys).await?;
}
} }
} }
} }
}).collect();
panic!("Each outer enum variant must be a tuple variant like `Alpaca(AlpacaCmd)`"); // Generate the final code
});
// Generate the final program
let expanded = quote! { let expanded = quote! {
use clap::Parser; use clap::Parser;
use std::io::Read; use std::io::Read;
@ -237,35 +231,19 @@ pub fn alpaca_cli(_attr: TokenStream, item: TokenStream) -> TokenStream {
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
let cmd = #top_enum_ident::parse(); let cmd = #top_enum_ident::parse();
match cmd { match cmd {
#(#outer_match_arms),* #(#match_arms),*
} }
Ok(()) Ok(())
} }
// Helper trait to unify async calls on sub-commands // Trait for dispatching API calls
trait ApiDispatch { pub trait ApiDispatch {
fn send_all( fn send_all(
&self, &self,
client: &awc::Client, client: &awc::Client,
keys: &std::collections::HashMap<String, crate::Keys>, keys: &std::collections::HashMap<String, crate::Keys>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), Box<dyn std::error::Error>>> + Send>>; ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), Box<dyn std::error::Error>>> + Send>>;
} }
// Implement ApiDispatch for every subcommand variant
#(impl ApiDispatch for #top_enum_ident {
fn send_all(
&self,
client: &awc::Client,
keys: &std::collections::HashMap<String, crate::Keys>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), Box<dyn std::error::Error>>> + Send>> {
Box::pin(async move {
match self {
#(#outer_match_arms),*
}
Ok(())
})
}
})*
}; };
TokenStream::from(expanded) TokenStream::from(expanded)