Skip to content
Snippets Groups Projects
Commit 63f1d0d5 authored by Simon Sapin's avatar Simon Sapin
Browse files

Share some more logic into the datasets module

parent b2a8e4ba
No related branches found
No related tags found
1 merge request!1Initial implementation
#[derive(serde::Serialize, serde::Deserialize)]
use crate::Megafine;
use serde::Deserialize;
use serde::Serialize;
use std::path::PathBuf;
#[derive(Serialize, Deserialize)]
pub(crate) struct DatasetsToml {
pub(crate) datasets: Vec<Dataset>,
}
#[derive(Serialize, Deserialize)]
pub(crate) struct Dataset {
pub(crate) name: String,
pub(crate) url: String,
pub(crate) sha256: String,
}
......@@ -2,5 +12,19 @@
pub(crate) struct Dataset {
pub(crate) name: String,
pub(crate) url: String,
pub(crate) sha256: String,
}
impl Dataset {
pub(crate) fn archive_name(&self) -> &str {
self.url.rsplit('/').next().unwrap()
}
pub(crate) fn cache_path(&self, megafine: &Megafine) -> PathBuf {
megafine
.cache_dir
.join("datasets")
.join(&self.sha256)
.join(self.archive_name())
}
}
use crate::datasets::Dataset;
use crate::datasets::DatasetsToml;
use crate::errors::Result;
use crate::utils::Tee;
use crate::Megafine;
......@@ -17,10 +18,5 @@
datasets_toml: PathBuf,
}
#[derive(serde::Deserialize)]
struct DatasetsToml {
datasets: Vec<Dataset>,
}
pub(crate) fn run(megafine: &Megafine, args: &Args) -> Result<()> {
let toml = fs::read(&args.datasets_toml)?;
......@@ -25,5 +21,5 @@
pub(crate) fn run(megafine: &Megafine, args: &Args) -> Result<()> {
let toml = fs::read(&args.datasets_toml)?;
let parsed: DatasetsToml = toml::from_slice(&toml)?;
let DatasetsToml { datasets } = toml::from_slice(&toml)?;
let client = reqwest::blocking::Client::new();
// TODO: write async code to download in parallel?
......@@ -28,13 +24,8 @@
let client = reqwest::blocking::Client::new();
// TODO: write async code to download in parallel?
for dataset in parsed.datasets {
let archive_name = dataset.url.rsplit('/').next().unwrap();
let cache_path = megafine
.cache_dir
.join("datasets")
.join(&dataset.sha256)
.join(&archive_name);
for dataset in datasets {
let cache_path = dataset.cache_path(megafine);
if cache_path.exists() {
// Trust our past self to have only renamed to that path
// a complete file with the correct SHA-256 hash.
......@@ -37,8 +28,8 @@
if cache_path.exists() {
// Trust our past self to have only renamed to that path
// a complete file with the correct SHA-256 hash.
println!("Already in cache: {}", style(&archive_name).bold());
println!("Already in cache: {}", style(dataset.archive_name()).bold());
continue;
}
println!("Downloading {}", style(&dataset.url).bold());
......
use crate::datasets::Dataset;
use crate::datasets::DatasetsToml;
use crate::errors::Result;
use crate::utils::subprocess_output;
use crate::utils::Tee;
......@@ -34,11 +35,6 @@
source_directory: PathBuf,
}
#[derive(serde::Serialize)]
struct NewDatasetToml {
datasets: [Dataset; 1],
}
pub(crate) fn run(megafine: &Megafine, args: &Args) -> Result<()> {
println!(
"New dataset from directory {}",
......@@ -130,11 +126,12 @@
println!("Compressed to {}", style(BinaryBytes(size)).bold());
println!("SHA-256: {}", style(&hash).bold());
let cache_path = megafine
.cache_dir
.join("datasets")
.join(&hash)
.join(&archive_name);
let dataset = Dataset {
url: format!("{}{}", DOWNLOAD_BASE_URL, archive_name),
name: dataset_name,
sha256: hash,
};
let cache_path = dataset.cache_path(megafine);
fs::create_dir_all(cache_path.parent().unwrap())?;
if args.overwrite {
tmp_file.persist(&cache_path)?;
......@@ -154,12 +151,8 @@
Err("")?
}
let to_toml = NewDatasetToml {
datasets: [Dataset {
url: format!("{}{}", DOWNLOAD_BASE_URL, archive_name),
name: dataset_name,
sha256: hash,
}],
let to_toml = DatasetsToml {
datasets: vec![dataset],
};
let toml_path = megafine.datasets.join("new.toml");
let mut toml_file = fs::OpenOptions::new()
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment