From 92c041fd7cbedc89ecb886140bc394b28ae4c518 Mon Sep 17 00:00:00 2001 From: mat Date: Wed, 3 Jan 2024 22:21:50 -0600 Subject: [PATCH] clean up engine declarations with macros --- src/engines/macros.rs | 124 ++++++++++++++++ src/engines/mod.rs | 180 +++++++----------------- src/engines/postsearch/docs_rs.rs | 7 +- src/engines/postsearch/github.rs | 2 +- src/engines/postsearch/stackexchange.rs | 2 +- 5 files changed, 184 insertions(+), 131 deletions(-) create mode 100644 src/engines/macros.rs diff --git a/src/engines/macros.rs b/src/engines/macros.rs new file mode 100644 index 0000000..ee67051 --- /dev/null +++ b/src/engines/macros.rs @@ -0,0 +1,124 @@ +#[macro_export] +macro_rules! engines { + ($($engine:ident = $id:expr),* $(,)?) => { + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] + pub enum Engine { + $($engine,)* + } + + impl Engine { + pub fn all() -> &'static [Engine] { + &[$(Engine::$engine,)*] + } + + pub fn id(&self) -> &'static str { + match self { + $(Engine::$engine => $id,)* + } + } + } + }; +} + +#[macro_export] +macro_rules! engine_weights { + ($($engine:ident = $weight:expr),* $(,)?) => { + impl Engine { + pub fn weight(&self) -> f64 { + match self { + $(Engine::$engine => $weight,)* + _ => 1., + } + } + } + }; +} + +#[macro_export] +macro_rules! engine_parse_response { + ($res:ident, $module:ident::$engine_id:ident::None) => { + None + }; + ($res:ident, $module:ident::$engine_id:ident::$parse_response:ident) => { + Some($module::$engine_id::$parse_response($res.into())) + }; +} + +#[macro_export] +macro_rules! engine_requests { + ($($engine:ident => $module:ident::$engine_id:ident::$request:ident, $parse_response:ident),* $(,)?) => { + impl Engine { + pub fn request(&self, query: &SearchQuery) -> RequestResponse { + #[allow(clippy::useless_conversion)] + match self { + $( + Engine::$engine => $module::$engine_id::$request(query).into(), + )* + _ => RequestResponse::None, + } + } + + pub fn parse_response(&self, res: &HttpResponse) -> eyre::Result { + #[allow(clippy::useless_conversion)] + match self { + $( + Engine::$engine => $crate::engine_parse_response! { res, $module::$engine_id::$parse_response } + .ok_or_else(|| eyre::eyre!("engine {self:?} can't parse response"))?, + )* + _ => eyre::bail!("engine {self:?} can't parse response"), + } + } + } + }; +} + +#[macro_export] +macro_rules! engine_autocomplete_requests { + ($($engine:ident => $module:ident::$engine_id:ident::$request:ident, $parse_response:ident),* $(,)?) => { + impl Engine { + pub fn request_autocomplete(&self, query: &str) -> Option { + match self { + $( + Engine::$engine => Some($module::$engine_id::$request(query).into()), + )* + _ => None, + } + } + + pub fn parse_autocomplete_response(&self, body: &str) -> eyre::Result> { + match self { + $( + Engine::$engine => $crate::engine_parse_response! { body, $module::$engine_id::$parse_response } + .ok_or_else(|| eyre::eyre!("engine {self:?} can't parse autocomplete response"))?, + )* + _ => eyre::bail!("engine {self:?} can't parse autocomplete response"), + } + } + } + }; +} + +#[macro_export] +macro_rules! engine_postsearch_requests { + ($($engine:ident => $module:ident::$engine_id:ident::$request:ident, $parse_response:ident),* $(,)?) => { + impl Engine { + pub fn postsearch_request(&self, response: &Response) -> Option { + match self { + $( + Engine::$engine => $module::$engine_id::$request(response), + )* + _ => None, + } + } + + pub fn postsearch_parse_response(&self, res: &HttpResponse) -> Option { + match self { + $( + Engine::$engine => $crate::engine_parse_response! { res, $module::$engine_id::$parse_response }?, + )* + _ => None, + } + } + } + }; +} diff --git a/src/engines/mod.rs b/src/engines/mod.rs index 20efe14..c285ea4 100644 --- a/src/engines/mod.rs +++ b/src/engines/mod.rs @@ -11,141 +11,64 @@ use std::{ use futures::future::join_all; use reqwest::header::HeaderMap; use tokio::sync::mpsc; -use url::Url; + +mod macros; +use crate::{ + engine_autocomplete_requests, engine_postsearch_requests, engine_requests, engine_weights, + engines, +}; pub mod answer; pub mod postsearch; pub mod search; -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub enum Engine { +engines! { // search - Google, - Bing, - Brave, - Marginalia, + Google = "google", + Bing = "bing", + Brave = "brave", + Marginalia = "marginalia", // answer - Useragent, - Ip, - Calc, - Wikipedia, - Dictionary, + Useragent = "useragent", + Ip = "ip", + Calc = "calc", + Wikipedia = "wikipedia", + Dictionary = "dictionary", // post-search - StackExchange, - GitHub, - DocsRs, + StackExchange = "stackexchange", + GitHub = "github", + DocsRs = "docs.rs", } -impl Engine { - pub fn all() -> &'static [Engine] { - &[ - Engine::Google, - Engine::Bing, - Engine::Brave, - Engine::Marginalia, - // answer - Engine::Useragent, - Engine::Ip, - Engine::Calc, - Engine::Wikipedia, - Engine::Dictionary, - // post-search - Engine::StackExchange, - Engine::GitHub, - Engine::DocsRs, - ] - } +engine_weights! { + Google = 1.05, + Bing = 1.0, + Brave = 1.25, + Marginalia = 0.15, + // defaults to 1.0 +} - pub fn id(&self) -> &'static str { - match self { - Engine::Google => "google", - Engine::Bing => "bing", - Engine::Brave => "brave", - Engine::Marginalia => "marginalia", - // answer - Engine::Useragent => "useragent", - Engine::Ip => "ip", - Engine::Calc => "calc", - Engine::Wikipedia => "wikipedia", - Engine::Dictionary => "dictionary", - // post-search - Engine::StackExchange => "stackexchange", - Engine::GitHub => "github", - Engine::DocsRs => "docs.rs", - } - } +engine_requests! { + Google => search::google::request, parse_response, + Bing => search::bing::request, parse_response, + Brave => search::brave::request, parse_response, + Marginalia => search::marginalia::request, parse_response, + Useragent => answer::useragent::request, None, + Ip => answer::ip::request, None, + Calc => answer::calc::request, None, + Wikipedia => answer::wikipedia::request, parse_response, + Dictionary => answer::dictionary::request, parse_response, +} - pub fn weight(&self) -> f64 { - match self { - Engine::Google => 1.05, - Engine::Bing => 1., - Engine::Brave => 1.25, - Engine::Marginalia => 0.15, - _ => 1., - } - } +engine_autocomplete_requests! { + Google => search::google::request_autocomplete, parse_autocomplete_response, + Calc => answer::calc::request_autocomplete, None, +} - pub fn request(&self, query: &SearchQuery) -> RequestResponse { - #[allow(clippy::useless_conversion)] - match self { - Engine::Google => search::google::request(query).into(), - Engine::Bing => search::bing::request(query).into(), - Engine::Brave => search::brave::request(query).into(), - Engine::Marginalia => search::marginalia::request(query).into(), - Engine::Useragent => answer::useragent::request(query).into(), - Engine::Ip => answer::ip::request(query).into(), - Engine::Calc => answer::calc::request(query).into(), - Engine::Wikipedia => answer::wikipedia::request(query).into(), - Engine::Dictionary => answer::dictionary::request(query).into(), - _ => RequestResponse::None, - } - } - - pub fn parse_response(&self, res: &HttpResponse) -> eyre::Result { - #[allow(clippy::useless_conversion)] - match self { - Engine::Google => search::google::parse_response(res.into()), - Engine::Bing => search::bing::parse_response(res.into()), - Engine::Brave => search::brave::parse_response(res.into()), - Engine::Marginalia => search::marginalia::parse_response(res.into()), - Engine::Wikipedia => answer::wikipedia::parse_response(res.into()), - Engine::Dictionary => answer::dictionary::parse_response(res.into()), - _ => eyre::bail!("engine {self:?} can't parse response"), - } - } - - pub fn request_autocomplete(&self, query: &str) -> Option { - match self { - Engine::Google => Some(search::google::request_autocomplete(query).into()), - Engine::Calc => Some(answer::calc::request_autocomplete(query).into()), - _ => None, - } - } - - pub fn parse_autocomplete_response(&self, body: &str) -> eyre::Result> { - match self { - Engine::Google => search::google::parse_autocomplete_response(body), - _ => eyre::bail!("engine {self:?} can't parse autocomplete response"), - } - } - - pub fn postsearch_request(&self, response: &Response) -> Option { - match self { - Engine::StackExchange => postsearch::stackexchange::request(response), - Engine::GitHub => postsearch::github::request(response), - Engine::DocsRs => postsearch::docs_rs::request(response), - _ => None, - } - } - - pub fn postsearch_parse_response(&self, body: &str, url: Url) -> Option { - match self { - Engine::StackExchange => postsearch::stackexchange::parse_response(body, url), - Engine::GitHub => postsearch::github::parse_response(body, url), - Engine::DocsRs => postsearch::docs_rs::parse_response(body, url), - _ => None, - } - } +engine_postsearch_requests! { + StackExchange => postsearch::stackexchange::request, parse_response, + GitHub => postsearch::github::request, parse_response, + DocsRs => postsearch::docs_rs::request, parse_response, } impl fmt::Display for Engine { @@ -393,10 +316,15 @@ pub async fn search_with_engines( if let Some(request) = engine.postsearch_request(&response) { postsearch_requests.push(async { let response = match request.send().await { - Ok(res) => { - let url = res.url().clone(); - let body = res.text().await?; - engine.postsearch_parse_response(&body, url) + Ok(mut res) => { + let mut body_bytes = Vec::new(); + while let Some(chunk) = res.chunk().await? { + body_bytes.extend_from_slice(&chunk); + } + let body = String::from_utf8_lossy(&body_bytes).to_string(); + + let http_response = HttpResponse { res, body }; + engine.postsearch_parse_response(&http_response) } Err(e) => { eprintln!("postsearch request error: {}", e); diff --git a/src/engines/postsearch/docs_rs.rs b/src/engines/postsearch/docs_rs.rs index fd915e8..34d41f1 100644 --- a/src/engines/postsearch/docs_rs.rs +++ b/src/engines/postsearch/docs_rs.rs @@ -1,7 +1,6 @@ use scraper::{Html, Selector}; -use url::Url; -use crate::engines::{Response, CLIENT}; +use crate::engines::{HttpResponse, Response, CLIENT}; pub fn request(response: &Response) -> Option { for search_result in response.search_results.iter().take(8) { @@ -13,7 +12,9 @@ pub fn request(response: &Response) -> Option { None } -pub fn parse_response(body: &str, url: Url) -> Option { +pub fn parse_response(HttpResponse { res, body }: &HttpResponse) -> Option { + let url = res.url().clone(); + let dom = Html::parse_document(body); let version = dom diff --git a/src/engines/postsearch/github.rs b/src/engines/postsearch/github.rs index 7b182ba..e9d144d 100644 --- a/src/engines/postsearch/github.rs +++ b/src/engines/postsearch/github.rs @@ -13,7 +13,7 @@ pub fn request(response: &Response) -> Option { None } -pub fn parse_response(body: &str, _url: Url) -> Option { +pub fn parse_response(body: &str) -> Option { let dom = Html::parse_document(body); let url_relative = dom diff --git a/src/engines/postsearch/stackexchange.rs b/src/engines/postsearch/stackexchange.rs index 71e0330..fc04854 100644 --- a/src/engines/postsearch/stackexchange.rs +++ b/src/engines/postsearch/stackexchange.rs @@ -15,7 +15,7 @@ pub fn request(response: &Response) -> Option { None } -pub fn parse_response(body: &str, _url: Url) -> Option { +pub fn parse_response(body: &str) -> Option { let dom = Html::parse_document(body); let title = dom