Compare commits
4 Commits
3d97489712
...
main
Author | SHA1 | Date | |
---|---|---|---|
83d5d85962 | |||
41f85afa6c | |||
0ba23fad0e | |||
37fa009fdd |
235
src/lib.rs
235
src/lib.rs
@ -3,13 +3,13 @@ use quote::quote;
|
|||||||
use syn::{parse_macro_input, Lit, ItemEnum, DeriveInput, Fields, Data};
|
use syn::{parse_macro_input, Lit, ItemEnum, DeriveInput, Fields, Data};
|
||||||
use quote::format_ident;
|
use quote::format_ident;
|
||||||
|
|
||||||
#[proc_macro_derive(HttpRequest, attributes(http_get, http_response))]
|
#[proc_macro_derive(HttpRequest, attributes(http_get, http_response, http_error_type))]
|
||||||
pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
|
pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
|
||||||
let input = parse_macro_input!(input as DeriveInput);
|
let input = parse_macro_input!(input as DeriveInput);
|
||||||
let query_name = &input.ident;
|
let query_name = &input.ident;
|
||||||
let query_name_str = query_name.to_string();
|
let query_name_str = query_name.to_string();
|
||||||
|
|
||||||
// Parse optional #[http_response = "..."] attribute via parse_nested_meta
|
// Parse optional #[http_response = "..."] attribute
|
||||||
let mut response_name_opt: Option<String> = None;
|
let mut response_name_opt: Option<String> = None;
|
||||||
for attr in &input.attrs {
|
for attr in &input.attrs {
|
||||||
if attr.path().is_ident("http_response") {
|
if attr.path().is_ident("http_response") {
|
||||||
@ -37,6 +37,22 @@ pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
|
|||||||
};
|
};
|
||||||
let response_name = format_ident!("{}", response_name_str);
|
let response_name = format_ident!("{}", response_name_str);
|
||||||
|
|
||||||
|
// Parse optional #[http_error_type = "..."] attribute (default to `current_mod::E`)
|
||||||
|
let mut error_type = syn::Path::from(syn::Ident::new("E", proc_macro2::Span::call_site()));
|
||||||
|
for attr in &input.attrs {
|
||||||
|
if attr.path().is_ident("http_error_type") {
|
||||||
|
attr.parse_nested_meta(|meta| {
|
||||||
|
if meta.path.is_ident("http_error_type") {
|
||||||
|
let lit: Lit = meta.value()?.parse()?;
|
||||||
|
if let Lit::Str(litstr) = lit {
|
||||||
|
error_type = syn::parse_str(&litstr.value()).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}).unwrap_or_else(|e| panic!("Error parsing http_error_type attribute: {}", e));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Extract base URL from #[http_get(url = "...")]
|
// Extract base URL from #[http_get(url = "...")]
|
||||||
let mut base_url = None;
|
let mut base_url = None;
|
||||||
for attr in &input.attrs {
|
for attr in &input.attrs {
|
||||||
@ -85,7 +101,7 @@ pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
|
|||||||
async fn send(
|
async fn send(
|
||||||
&self,
|
&self,
|
||||||
headers: Option<Vec<(&str, &str)>>,
|
headers: Option<Vec<(&str, &str)>>,
|
||||||
) -> Result<Self::R, Box<dyn std::error::Error + Send + Sync>> {
|
) -> Result<Self::R, #error_type> { // Use the error type here
|
||||||
use awc::Client;
|
use awc::Client;
|
||||||
use urlencoding::encode;
|
use urlencoding::encode;
|
||||||
|
|
||||||
@ -131,119 +147,158 @@ 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 enum_name = &input_enum.ident;
|
let top_enum_ident = &input_enum.ident;
|
||||||
let variants = &input_enum.variants;
|
let top_variants = &input_enum.variants;
|
||||||
|
|
||||||
// Match arms for regular command variants (Single, Asset, etc.)
|
// Build outer match arms
|
||||||
let regular_arms = variants.iter().filter_map(|v| {
|
let match_arms: Vec<_> = top_variants.iter().map(|variant| {
|
||||||
let v_name = &v.ident;
|
let variant_ident = &variant.ident;
|
||||||
|
|
||||||
// Skip Bulk variant — we handle it separately
|
// Expecting tuple variants like Alpaca(AlpacaCmd)
|
||||||
if v_name == "Bulk" {
|
let inner_type = match &variant.fields {
|
||||||
return None;
|
Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
|
||||||
|
match &fields.unnamed.first().unwrap().ty {
|
||||||
|
syn::Type::Path(p) => p.path.segments.last().unwrap().ident.clone(),
|
||||||
|
_ => panic!("Expected tuple variant with a type path"),
|
||||||
}
|
}
|
||||||
|
|
||||||
Some(quote! {
|
|
||||||
#enum_name::#v_name(req) => {
|
|
||||||
let res = req.send(client.clone(), &api_key).await?;
|
|
||||||
let body = res.body().await?;
|
|
||||||
println!("{}", std::str::from_utf8(&body)?);
|
|
||||||
}
|
}
|
||||||
})
|
_ => panic!("Each variant must be a tuple variant like `Alpaca(AlpacaCmd)`"),
|
||||||
});
|
};
|
||||||
|
|
||||||
let expanded = quote! {
|
quote! {
|
||||||
#[derive(structopt::StructOpt, Debug)]
|
#top_enum_ident::#variant_ident(inner) => {
|
||||||
#input_enum
|
match inner {
|
||||||
|
#inner_type::Bulk { input } => {
|
||||||
#[tokio::main]
|
let mut reader: Box<dyn std::io::Read> = match input {
|
||||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
Some(path) => Box::new(std::fs::File::open(path)?),
|
||||||
use structopt::StructOpt;
|
|
||||||
use std::fs::File;
|
|
||||||
use std::io::{BufReader, Read};
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::thread;
|
|
||||||
|
|
||||||
const THREADS: usize = 4;
|
|
||||||
|
|
||||||
// Initialize shared HTTP client and API key
|
|
||||||
let client = Arc::new(awc::Client::default());
|
|
||||||
let api_key = std::env::var("APCA_API_KEY_ID")?;
|
|
||||||
let cmd = #enum_name::from_args();
|
|
||||||
|
|
||||||
match cmd {
|
|
||||||
#(#regular_arms)*
|
|
||||||
|
|
||||||
#enum_name::Bulk { input } => {
|
|
||||||
// Choose input source: file or stdin
|
|
||||||
let mut reader: Box<dyn Read> = match input {
|
|
||||||
Some(path) => Box::new(File::open(path)?),
|
|
||||||
None => Box::new(std::io::stdin()),
|
None => Box::new(std::io::stdin()),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Read input JSON into buffer
|
|
||||||
let mut buf = String::new();
|
let mut buf = String::new();
|
||||||
reader.read_to_string(&mut buf)?;
|
reader.read_to_string(&mut buf)?;
|
||||||
|
let queries: Vec<#inner_type> = serde_json::from_str(&buf)?;
|
||||||
|
|
||||||
// Deserialize into Vec<Query>
|
use std::sync::Arc;
|
||||||
let queries: Vec<Query> = serde_json::from_str(&buf)?;
|
let client = Arc::new(awc::Client::default());
|
||||||
|
let keys = Arc::new(crate::load_api_keys()?);
|
||||||
|
|
||||||
|
const THREADS: usize = 4;
|
||||||
let total = queries.len();
|
let total = queries.len();
|
||||||
|
let per_thread = std::cmp::max(1, total / THREADS);
|
||||||
if total == 0 {
|
|
||||||
eprintln!("No queries provided.");
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
let shared_queries = Arc::new(queries);
|
let shared_queries = Arc::new(queries);
|
||||||
let shared_key = Arc::new(api_key);
|
|
||||||
|
|
||||||
let per_thread = total / THREADS;
|
let mut handles = Vec::new();
|
||||||
|
|
||||||
let mut handles = Vec::with_capacity(THREADS);
|
|
||||||
for i in 0..THREADS {
|
for i in 0..THREADS {
|
||||||
let queries_clone = Arc::clone(&shared_queries);
|
let queries = Arc::clone(&shared_queries);
|
||||||
let client_clone = Arc::clone(&client);
|
let client = Arc::clone(&client);
|
||||||
let key_clone = Arc::clone(&shared_key);
|
let keys = Arc::clone(&keys);
|
||||||
|
let start = i * per_thread;
|
||||||
|
let end = if i == THREADS - 1 { total } else { start + per_thread };
|
||||||
|
|
||||||
let start_index = i * per_thread;
|
let handle = std::thread::spawn(move || {
|
||||||
let end_index = if i == THREADS - 1 {
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||||
total // Last thread gets the remainder
|
for q in &queries[start..end] {
|
||||||
} else {
|
rt.block_on(q.send_all(&client, &keys)).unwrap();
|
||||||
start_index + per_thread
|
|
||||||
};
|
|
||||||
|
|
||||||
let handle = thread::spawn(move || {
|
|
||||||
let rt = tokio::runtime::Runtime::new().expect("Failed to create runtime");
|
|
||||||
for idx in start_index..end_index {
|
|
||||||
let query = &queries_clone[idx];
|
|
||||||
let send_result = rt.block_on(query.send(client_clone.clone(), &key_clone));
|
|
||||||
|
|
||||||
match send_result {
|
|
||||||
Ok(response) => {
|
|
||||||
let body_result = rt.block_on(response.body());
|
|
||||||
match body_result {
|
|
||||||
Ok(body) => println!("{}", String::from_utf8_lossy(&body)),
|
|
||||||
Err(e) => eprintln!("Error reading response body: {:?}", e),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
eprintln!("Request failed: {:?}", e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
handles.push(handle);
|
handles.push(handle);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Wait for all threads to complete
|
for h in handles {
|
||||||
for handle in handles {
|
h.join().expect("Thread panicked");
|
||||||
handle.join().expect("A thread panicked");
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
other => {
|
||||||
|
let client = awc::Client::default();
|
||||||
|
let keys = crate::load_api_keys()?;
|
||||||
|
other.send_all(&client, &keys).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}).collect();
|
||||||
|
|
||||||
|
// Generate the final code
|
||||||
|
let expanded = quote! {
|
||||||
|
use clap::Parser;
|
||||||
|
use std::io::Read;
|
||||||
|
|
||||||
|
#[derive(clap::Parser, Debug)]
|
||||||
|
#input_enum
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
let cmd = #top_enum_ident::parse();
|
||||||
|
match cmd {
|
||||||
|
#(#match_arms),*
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
// Trait for dispatching API calls
|
||||||
|
pub trait ApiDispatch {
|
||||||
|
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>>;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
TokenStream::from(expanded)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[proc_macro_attribute]
|
||||||
|
pub fn api_dispatch(attr: TokenStream, item: TokenStream) -> TokenStream {
|
||||||
|
let input = parse_macro_input!(item as syn::ItemEnum);
|
||||||
|
let enum_ident = &input.ident;
|
||||||
|
|
||||||
|
// Parse attribute input: input = "MyQuery"
|
||||||
|
let meta_args = attr.to_string();
|
||||||
|
let input_type: syn::Ident = {
|
||||||
|
let cleaned = meta_args.trim().replace("input", "").replace('=', "").replace('"', "").trim().to_string();
|
||||||
|
syn::Ident::new(&cleaned, proc_macro2::Span::call_site())
|
||||||
|
};
|
||||||
|
|
||||||
|
let expanded = quote! {
|
||||||
|
#input
|
||||||
|
|
||||||
|
impl ApiDispatch for #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 {
|
||||||
|
#enum_ident::Single { query } => {
|
||||||
|
let parsed: #input_type = serde_json::from_str(query)?;
|
||||||
|
parsed.send(client, keys).await?;
|
||||||
|
}
|
||||||
|
#enum_ident::Bulk { input } => {
|
||||||
|
let json = if let Some(raw) = input {
|
||||||
|
if std::path::Path::new(&raw).exists() {
|
||||||
|
std::fs::read_to_string(raw)?
|
||||||
|
} else {
|
||||||
|
raw.clone()
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
use std::io::Read;
|
||||||
|
let mut buf = String::new();
|
||||||
|
std::io::stdin().read_to_string(&mut buf)?;
|
||||||
|
buf
|
||||||
|
};
|
||||||
|
|
||||||
|
let items: Vec<#input_type> = serde_json::from_str(&json)?;
|
||||||
|
for item in items {
|
||||||
|
item.send(client, keys).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user