Skip to main content

paper_gap/
pwc.rs

1use crate::Result;
2use crate::arxiv::Paper;
3use serde::{Deserialize, Serialize};
4use std::time::Duration;
5
6const BASE_URL: &str = "https://paperswithcode.com/api/v1";
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9pub struct CodeLink {
10    pub source_provider: String,
11    pub paper_id: Option<String>,
12    pub repository_url: String,
13    pub repository_name: Option<String>,
14    pub is_official: Option<bool>,
15}
16
17#[derive(Debug, Deserialize)]
18struct List<T> {
19    results: Vec<T>,
20}
21
22#[derive(Debug, Deserialize)]
23struct PwcPaper {
24    id: String,
25}
26
27#[derive(Debug, Deserialize)]
28struct PwcRepository {
29    url: Option<String>,
30    name: Option<String>,
31    is_official: Option<bool>,
32}
33
34pub fn paper_search_url(paper: &Paper) -> String {
35    if let Some(arxiv_id) = paper.arxiv_id.as_deref() {
36        format!(
37            "{BASE_URL}/papers/?arxiv_id={}",
38            urlencoding::encode(trim_arxiv_version(arxiv_id))
39        )
40    } else {
41        format!("{BASE_URL}/papers/?q={}", urlencoding::encode(&paper.title))
42    }
43}
44
45pub fn paper_repositories_url(pwc_paper_id: &str) -> String {
46    format!(
47        "{BASE_URL}/papers/{}/repositories/",
48        urlencoding::encode(pwc_paper_id)
49    )
50}
51
52pub async fn code_links(paper: &Paper) -> Result<Vec<CodeLink>> {
53    let client = reqwest::Client::builder()
54        .timeout(Duration::from_secs(8))
55        .build()?;
56    let paper_response = client.get(paper_search_url(paper)).send().await?;
57    if paper_response.status() == reqwest::StatusCode::NOT_FOUND {
58        return Ok(Vec::new());
59    }
60    let paper_response = paper_response.error_for_status()?;
61    let papers = parse_paper_search(&paper_response.text().await?)?;
62    let Some(pwc_paper) = papers.first() else {
63        return Ok(Vec::new());
64    };
65
66    let repo_response = client
67        .get(paper_repositories_url(&pwc_paper.id))
68        .send()
69        .await?;
70    if repo_response.status() == reqwest::StatusCode::NOT_FOUND {
71        return Ok(Vec::new());
72    }
73    let repo_response = repo_response.error_for_status()?;
74    parse_repositories(&pwc_paper.id, &repo_response.text().await?)
75}
76
77fn parse_paper_search(json: &str) -> Result<Vec<PwcPaper>> {
78    Ok(serde_json::from_str::<List<PwcPaper>>(json)?.results)
79}
80
81pub fn parse_repositories(paper_id: &str, json: &str) -> Result<Vec<CodeLink>> {
82    let repos = serde_json::from_str::<List<PwcRepository>>(json)?;
83    Ok(repos
84        .results
85        .into_iter()
86        .filter_map(|repo| {
87            Some(CodeLink {
88                source_provider: "papers-with-code".into(),
89                paper_id: Some(paper_id.into()),
90                repository_url: repo.url?,
91                repository_name: repo.name,
92                is_official: repo.is_official,
93            })
94        })
95        .collect())
96}
97
98fn trim_arxiv_version(arxiv_id: &str) -> &str {
99    arxiv_id.rsplit_once('v').map_or(arxiv_id, |(id, version)| {
100        if version.chars().all(|c| c.is_ascii_digit()) {
101            id
102        } else {
103            arxiv_id
104        }
105    })
106}