Compare commits
38 Commits
999433cacf
...
main
Author | SHA1 | Date | |
---|---|---|---|
83d5d85962 | |||
41f85afa6c | |||
0ba23fad0e | |||
37fa009fdd | |||
3d97489712 | |||
1524769694 | |||
0a88b1d7a1 | |||
4b1fac6403 | |||
a9e59b6865 | |||
3e75cd00ff | |||
a096885f02 | |||
5b3641d930 | |||
981e662bd3 | |||
2c84a8bb18 | |||
f1a6786419 | |||
919c5eda3e | |||
8dcd080f5b | |||
9fc0ab1f0d | |||
5aa65e1a38 | |||
b3b813b092 | |||
dbba095e22 | |||
c89d5bf748 | |||
4504e1132b | |||
820f5445a3 | |||
e18eb7bb02 | |||
745415ba13 | |||
716a72fffd | |||
8264838b2c | |||
a6f8a94666 | |||
259c5562e4 | |||
d82f663c27 | |||
e96def50f3 | |||
533f718e57 | |||
f9ee10491d | |||
b11cecdff4 | |||
20413273aa | |||
ac0ace0b94 | |||
275d706845 |
12
Cargo.lock
generated
12
Cargo.lock
generated
@ -2,10 +2,22 @@
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.88"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e539d3fca749fcee5236ab05e93a52867dd549cc157c8cb7f99595f3cedffdb5"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "derive_http"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
|
@ -7,6 +7,7 @@ edition = "2024"
|
||||
proc-macro = true
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1.88"
|
||||
proc-macro2 = "1.0.95"
|
||||
quote = "1.0.40"
|
||||
syn = "2.0.104"
|
||||
syn = { version = "2.0.104", features = ["full"] }
|
||||
|
304
src/lib.rs
304
src/lib.rs
@ -1,75 +1,129 @@
|
||||
use proc_macro::TokenStream;
|
||||
use quote::quote;
|
||||
use syn::{parse_macro_input, DeriveInput, Data, Fields, Lit};
|
||||
use syn::{parse_macro_input, Lit, ItemEnum, DeriveInput, Fields, Data};
|
||||
use quote::format_ident;
|
||||
|
||||
#[proc_macro_derive(HttpRequest, attributes(http_get))]
|
||||
#[proc_macro_derive(HttpRequest, attributes(http_get, http_response, http_error_type))]
|
||||
pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
|
||||
let input = parse_macro_input!(input as DeriveInput);
|
||||
let name = &input.ident;
|
||||
let query_name = &input.ident;
|
||||
let query_name_str = query_name.to_string();
|
||||
|
||||
// Parse #[http_get(url = "...")] attribute
|
||||
// Parse optional #[http_response = "..."] attribute
|
||||
let mut response_name_opt: Option<String> = None;
|
||||
for attr in &input.attrs {
|
||||
if attr.path().is_ident("http_response") {
|
||||
attr.parse_nested_meta(|meta| {
|
||||
if meta.path.is_ident("http_response") {
|
||||
let lit: Lit = meta.value()?.parse()?;
|
||||
if let Lit::Str(litstr) = lit {
|
||||
response_name_opt = Some(litstr.value());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}).unwrap_or_else(|e| panic!("Error parsing http_response attribute: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// Determine response type name
|
||||
let response_name_str = if let Some(custom_resp) = response_name_opt {
|
||||
custom_resp
|
||||
} else if query_name_str == "Q" {
|
||||
"R".to_string()
|
||||
} else if query_name_str.ends_with('Q') {
|
||||
format!("{}R", &query_name_str[..query_name_str.len() - 1])
|
||||
} else {
|
||||
panic!("HttpRequest derive macro expects the type name to be 'Q' or end with 'Q', or specify #[http_response = \"...\"] to override");
|
||||
};
|
||||
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 = "...")]
|
||||
let mut base_url = None;
|
||||
for attr in &input.attrs {
|
||||
if attr.path().is_ident("http_get") {
|
||||
let _ = attr.parse_nested_meta(|meta| {
|
||||
attr.parse_nested_meta(|meta| {
|
||||
if meta.path.is_ident("url") {
|
||||
let value: Lit = meta.value()?.parse()?;
|
||||
if let Lit::Str(litstr) = value {
|
||||
let lit: Lit = meta.value()?.parse()?;
|
||||
if let Lit::Str(litstr) = lit {
|
||||
base_url = Some(litstr.value());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
});
|
||||
}).unwrap_or_else(|e| panic!("Error parsing http_get attribute: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
let base_url = base_url.expect("Missing #[http_get(url = \"...\")] attribute");
|
||||
let base_url_lit = syn::LitStr::new(&base_url, proc_macro2::Span::call_site());
|
||||
|
||||
let expanded = match &input.data {
|
||||
Data::Struct(data_struct) => {
|
||||
let fields = match &data_struct.fields {
|
||||
Fields::Named(named) => &named.named,
|
||||
_ => panic!("#[derive(HttpRequest)] only supports structs with named fields"),
|
||||
};
|
||||
|
||||
let mut query_param_code = Vec::new();
|
||||
for field in fields {
|
||||
let ident = field.ident.clone().unwrap();
|
||||
// Collect query parameters from fields named "lnk_p_*" (only for structs)
|
||||
let query_param_code = if let Data::Struct(data_struct) = &input.data {
|
||||
if let Fields::Named(fields_named) = &data_struct.fields {
|
||||
fields_named.named.iter().filter_map(|field| {
|
||||
let ident = field.ident.as_ref()?;
|
||||
let field_name = ident.to_string();
|
||||
if field_name.starts_with("lnk_p_") {
|
||||
let key = &field_name["lnk_p_".len()..];
|
||||
query_param_code.push(quote! {
|
||||
Some(quote! {
|
||||
query_params.push((#key.to_string(), self.#ident.to_string()));
|
||||
});
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}).collect::<Vec<_>>()
|
||||
} else {
|
||||
Vec::new()
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
quote! {
|
||||
impl #name {
|
||||
pub async fn send(
|
||||
let expanded = quote! {
|
||||
#[async_trait::async_trait]
|
||||
impl Queryable for #query_name {
|
||||
type R = #response_name;
|
||||
|
||||
async fn send(
|
||||
&self,
|
||||
client: std::sync::Arc<awc::Client>,
|
||||
headers: Option<Vec<(&str, &str)>>,
|
||||
api_key: Option<&str>,
|
||||
) -> Result<awc::ClientResponse, awc::error::SendRequestError> {
|
||||
) -> Result<Self::R, #error_type> { // Use the error type here
|
||||
use awc::Client;
|
||||
use urlencoding::encode;
|
||||
|
||||
let mut query_params: Vec<(String, String)> = Vec::new();
|
||||
#(#query_param_code)*
|
||||
|
||||
if let Some(key) = api_key {
|
||||
query_params.push(("api_key".to_string(), key.to_string()));
|
||||
}
|
||||
|
||||
let mut url = #base_url.to_string();
|
||||
let mut url = #base_url_lit.to_string();
|
||||
if !query_params.is_empty() {
|
||||
let query_parts: Vec<String> = query_params.iter()
|
||||
.map(|(k, v)| format!("{}={}", k, encode(v)))
|
||||
.collect();
|
||||
let mut query_string = String::new();
|
||||
let mut first = true;
|
||||
for (k, v) in &query_params {
|
||||
if !first {
|
||||
query_string.push('&');
|
||||
}
|
||||
first = false;
|
||||
query_string.push_str(&format!("{}={}", k, encode(v)));
|
||||
}
|
||||
url.push('?');
|
||||
url.push_str(&query_parts.join("&"));
|
||||
url.push_str(&query_string);
|
||||
}
|
||||
|
||||
let client = Client::default();
|
||||
let mut request = client.get(url);
|
||||
|
||||
if let Some(hdrs) = headers {
|
||||
@ -79,57 +133,171 @@ pub fn derive_http_get_request(input: TokenStream) -> TokenStream {
|
||||
}
|
||||
|
||||
let response = request.send().await?;
|
||||
Ok(response)
|
||||
}
|
||||
let bytes = response.body().await?;
|
||||
|
||||
let parsed: Self::R = serde_json::from_slice(&bytes)?;
|
||||
Ok(parsed)
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
TokenStream::from(expanded)
|
||||
}
|
||||
|
||||
Data::Enum(data_enum) => {
|
||||
let mut variant_arms = Vec::new();
|
||||
for variant in &data_enum.variants {
|
||||
let vname = &variant.ident;
|
||||
match &variant.fields {
|
||||
#[proc_macro_attribute]
|
||||
pub fn alpaca_cli(_attr: TokenStream, item: TokenStream) -> TokenStream {
|
||||
let input_enum = parse_macro_input!(item as ItemEnum);
|
||||
let top_enum_ident = &input_enum.ident;
|
||||
let top_variants = &input_enum.variants;
|
||||
|
||||
// Build outer match arms
|
||||
let match_arms: Vec<_> = top_variants.iter().map(|variant| {
|
||||
let variant_ident = &variant.ident;
|
||||
|
||||
// Expecting tuple variants like Alpaca(AlpacaCmd)
|
||||
let inner_type = match &variant.fields {
|
||||
Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
|
||||
variant_arms.push(quote! {
|
||||
#name::#vname(inner) => inner.send(client.clone(), headers.clone(), api_key).await,
|
||||
});
|
||||
}
|
||||
_ => panic!("#[derive(HttpRequest)] enum variants must have a single unnamed field"),
|
||||
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"),
|
||||
}
|
||||
}
|
||||
_ => panic!("Each variant must be a tuple variant like `Alpaca(AlpacaCmd)`"),
|
||||
};
|
||||
|
||||
quote! {
|
||||
impl #name {
|
||||
pub async fn send(
|
||||
&self,
|
||||
client: std::sync::Arc<awc::Client>,
|
||||
headers: Option<Vec<(&str, &str)>>,
|
||||
api_key: Option<&str>,
|
||||
) -> Result<awc::ClientResponse, awc::error::SendRequestError> {
|
||||
match self {
|
||||
#(#variant_arms)*
|
||||
}
|
||||
}
|
||||
}
|
||||
#top_enum_ident::#variant_ident(inner) => {
|
||||
match inner {
|
||||
#inner_type::Bulk { input } => {
|
||||
let mut reader: Box<dyn std::io::Read> = match input {
|
||||
Some(path) => Box::new(std::fs::File::open(path)?),
|
||||
None => Box::new(std::io::stdin()),
|
||||
};
|
||||
|
||||
let mut buf = String::new();
|
||||
reader.read_to_string(&mut buf)?;
|
||||
let queries: Vec<#inner_type> = serde_json::from_str(&buf)?;
|
||||
|
||||
use std::sync::Arc;
|
||||
let client = Arc::new(awc::Client::default());
|
||||
let keys = Arc::new(crate::load_api_keys()?);
|
||||
|
||||
const THREADS: usize = 4;
|
||||
let total = queries.len();
|
||||
let per_thread = std::cmp::max(1, total / THREADS);
|
||||
let shared_queries = Arc::new(queries);
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for i in 0..THREADS {
|
||||
let queries = Arc::clone(&shared_queries);
|
||||
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 };
|
||||
|
||||
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();
|
||||
}
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
_ => panic!("#[derive(HttpRequest)] only supports structs and enums"),
|
||||
for h in handles {
|
||||
h.join().expect("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(())
|
||||
}
|
||||
|
||||
// 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_derive(HttpResponse)]
|
||||
pub fn derive_http_response(input: TokenStream) -> TokenStream {
|
||||
let input = parse_macro_input!(input as DeriveInput);
|
||||
let name = &input.ident;
|
||||
#[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! {
|
||||
impl #name {
|
||||
pub async fn receive(resp: awc::ClientResponse) -> Result<Self, awc::error::JsonPayloadError> {
|
||||
resp.json().await
|
||||
#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