use crate::config::{load_config, PackConfig, PathsConfig};
use std::io::Write;
use std::path::{Path, PathBuf};
fn add_dir_to_zip<W: std::io::Write + std::io::Seek>(
zip: &mut zip::ZipWriter<W>,
base: &Path,
dir: &Path,
opts: zip::write::SimpleFileOptions,
) -> std::io::Result<()> {
for entry in std::fs::read_dir(dir)? {
let path = entry?.path();
let rel = path.strip_prefix(base).unwrap().to_string_lossy().replace('\\', "/");
if path.is_dir() {
zip.add_directory(format!("{rel}/"), opts)
.map_err(|e| std::io::Error::other(e.to_string()))?;
add_dir_to_zip(zip, base, &path, opts)?;
} else {
zip.start_file(rel, opts).map_err(|e| std::io::Error::other(e.to_string()))?;
let bytes = std::fs::read(&path)?;
zip.write_all(&bytes)?;
}
}
Ok(())
}
pub fn zip_dir(src_dir: &Path, zip_path: &Path) -> std::io::Result<()> {
if let Some(parent) = zip_path.parent() {
std::fs::create_dir_all(parent)?;
}
let file = std::fs::File::create(zip_path)?;
let mut zip = zip::ZipWriter::new(file);
let opts = zip::write::SimpleFileOptions::default()
.compression_method(zip::CompressionMethod::Deflated);
let base = src_dir.parent().unwrap_or(src_dir);
add_dir_to_zip(&mut zip, base, src_dir, opts)?;
zip.finish().map_err(|e| std::io::Error::other(e.to_string()))?;
Ok(())
}
pub fn run_pack(repo_root: &Path) -> std::io::Result<()> {
let cfg = load_config(&repo_root.join(crate::config::CONFIG_FILE));
let pack = cfg.pack.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"no [pack] table in kibble.toml — see the Kaggle runbook",
)
})?;
let bundle = stage_bundle(repo_root, &cfg.paths, &pack)?;
let zip_path = repo_root.join(&pack.zip);
zip_dir(&bundle, &zip_path)?;
println!("Packed bundle: {}", bundle.display());
println!("Zip: {}", zip_path.display());
println!("Upload: kaggle datasets version -p {} -m \"refresh\"", repo_root.join(&pack.staging).display());
println!("Then run the training notebook on Kaggle and download the adapter.");
Ok(())
}
fn copy_recursive(src: &Path, dest: &Path) -> std::io::Result<()> {
if src.is_dir() {
std::fs::create_dir_all(dest)?;
for entry in std::fs::read_dir(src)? {
let entry = entry?;
copy_recursive(&entry.path(), &dest.join(entry.file_name()))?;
}
} else {
if let Some(parent) = dest.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::copy(src, dest)?;
}
Ok(())
}
pub fn stage_bundle(
repo_root: &Path,
paths: &PathsConfig,
pack: &PackConfig,
) -> std::io::Result<PathBuf> {
let dataset_dir = repo_root.join(&paths.dataset_dir);
let staging_parent = repo_root.join(&pack.staging);
let bundle = staging_parent.join(&pack.name);
let _ = std::fs::remove_dir_all(&bundle);
std::fs::create_dir_all(bundle.join("full"))?;
for split in &pack.splits {
let src = dataset_dir.join(format!("{split}.jsonl"));
if !src.is_file() {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("missing split {} — run `kibble build` first", src.display()),
));
}
std::fs::copy(&src, bundle.join("full").join(format!("{split}.jsonl")))?;
}
let stats = dataset_dir.join("stats.json");
if stats.is_file() {
std::fs::copy(&stats, bundle.join("full/stats.json"))?;
}
for inc in &pack.include {
let src = repo_root.join(&inc.src);
if !src.exists() {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("pack include missing: {}", src.display()),
));
}
copy_recursive(&src, &bundle.join(&inc.dest))?;
}
let readme = pack.readme.clone().unwrap_or_else(|| "KIBBLE dataset bundle.".to_string());
std::fs::write(bundle.join("README.txt"), readme)?;
let meta = serde_json::json!({
"title": pack.metadata.title.clone().unwrap_or_default(),
"id": pack.metadata.id.clone().unwrap_or_default(),
"licenses": [{ "name": pack.metadata.license.clone().unwrap_or_else(|| "other".to_string()) }],
});
std::fs::write(
staging_parent.join("dataset-metadata.json"),
serde_json::to_string_pretty(&meta)?,
)?;
Ok(bundle)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{PackInclude, PackMetadata};
use std::fs;
fn paths() -> PathsConfig {
PathsConfig { data_root: "data".into(), dataset_dir: "data/datasets/unsloth".into() }
}
#[test]
fn stages_bundle_with_splits_includes_and_metadata() {
let root = std::env::temp_dir().join(format!("kibble_pack_{}", std::process::id()));
let _ = fs::remove_dir_all(&root);
fs::create_dir_all(root.join("data/datasets/unsloth")).unwrap();
fs::write(root.join("data/datasets/unsloth/train.jsonl"), "{\"x\":1}\n").unwrap();
fs::write(root.join("data/datasets/unsloth/valid.jsonl"), "{\"x\":2}\n").unwrap();
fs::write(root.join("data/datasets/unsloth/stats.json"), "{}").unwrap();
fs::create_dir_all(root.join("kaggle")).unwrap();
fs::write(root.join("kaggle/train_kaggle.py"), "print('hi')").unwrap();
let pack = PackConfig {
name: "style-lora".into(),
staging: "kaggle/dataset-staging".into(),
splits: vec!["train".into(), "valid".into()],
zip: "kaggle/out.zip".into(),
readme: Some("hello readme".into()),
include: vec![PackInclude { src: "kaggle/train_kaggle.py".into(), dest: "scripts/train_kaggle.py".into() }],
metadata: PackMetadata { id: Some("u/d".into()), title: Some("D".into()), license: Some("other".into()) },
};
let dir = stage_bundle(&root, &paths(), &pack).unwrap();
assert!(dir.ends_with("kaggle/dataset-staging/style-lora"));
assert!(dir.join("full/train.jsonl").is_file());
assert!(dir.join("full/valid.jsonl").is_file());
assert!(dir.join("full/stats.json").is_file());
assert!(dir.join("scripts/train_kaggle.py").is_file());
assert_eq!(fs::read_to_string(dir.join("README.txt")).unwrap(), "hello readme");
let meta = fs::read_to_string(root.join("kaggle/dataset-staging/dataset-metadata.json")).unwrap();
assert!(meta.contains("\"id\""));
assert!(meta.contains("u/d"));
assert!(meta.contains("\"licenses\""));
}
#[test]
fn errors_on_missing_split() {
let root = std::env::temp_dir().join(format!("kibble_pack_miss_{}", std::process::id()));
let _ = fs::remove_dir_all(&root);
fs::create_dir_all(root.join("data/datasets/unsloth")).unwrap();
let pack = PackConfig {
name: "b".into(), staging: "stg".into(), splits: vec!["train".into()],
zip: "z.zip".into(), readme: None, include: vec![], metadata: Default::default(),
};
assert!(stage_bundle(&root, &paths(), &pack).is_err());
}
#[test]
fn zip_dir_creates_archive_with_entries() {
use std::fs;
let root = std::env::temp_dir().join(format!("kibble_zip_{}", std::process::id()));
let _ = fs::remove_dir_all(&root);
let src = root.join("bundle");
fs::create_dir_all(src.join("full")).unwrap();
fs::write(src.join("full/train.jsonl"), "{}\n").unwrap();
fs::write(src.join("README.txt"), "hi").unwrap();
let zip_path = root.join("out.zip");
super::zip_dir(&src, &zip_path).unwrap();
assert!(zip_path.is_file());
let file = fs::File::open(&zip_path).unwrap();
let mut archive = zip::ZipArchive::new(file).unwrap();
let names: Vec<String> = (0..archive.len())
.map(|i| archive.by_index(i).unwrap().name().to_string())
.collect();
assert!(names.iter().any(|n| n.ends_with("full/train.jsonl")));
assert!(names.iter().any(|n| n.ends_with("README.txt")));
}
#[test]
fn errors_on_missing_include() {
let root = std::env::temp_dir().join(format!("kibble_pack_inc_{}", std::process::id()));
let _ = fs::remove_dir_all(&root);
fs::create_dir_all(root.join("data/datasets/unsloth")).unwrap();
fs::write(root.join("data/datasets/unsloth/train.jsonl"), "{}\n").unwrap();
let pack = PackConfig {
name: "b".into(), staging: "stg".into(), splits: vec!["train".into()],
zip: "z.zip".into(), readme: None,
include: vec![PackInclude { src: "nope.py".into(), dest: "scripts/nope.py".into() }],
metadata: Default::default(),
};
assert!(stage_bundle(&root, &paths(), &pack).is_err());
}
}