use crate::constants::buffer;
use anyhow::{Context, Result};
use serde_json::Value;
use std::fs::File;
use std::io::{BufRead, BufReader, Read, Write};
use std::path::Path;
pub fn read_jsonl<P: AsRef<Path>>(path: P) -> Result<Vec<Value>> {
let file = File::open(path.as_ref())
.with_context(|| format!("Failed to open file: {}", path.as_ref().display()))?;
let file_size = file.metadata().ok().map(|m| m.len() as usize).unwrap_or(0);
let estimated_lines = if file_size > 0 {
file_size / buffer::AVG_JSONL_LINE_SIZE
} else {
10 };
let mut results = Vec::with_capacity(estimated_lines);
let reader = BufReader::with_capacity(buffer::FILE_READ_BUFFER, file);
for (index, line) in reader.lines().enumerate() {
let line = line.with_context(|| format!("Failed to read line {}", index + 1))?;
if line.trim().is_empty() {
continue;
}
let obj: Value = serde_json::from_str(&line)
.with_context(|| format!("Failed to parse JSON at line {}", index + 1))?;
results.push(obj);
}
results.shrink_to_fit();
Ok(results)
}
pub fn write_json_atomic<T, P>(path: P, value: &T) -> Result<()>
where
T: serde::Serialize,
P: AsRef<Path>,
{
write_json_atomic_inner(path.as_ref(), value, false)
}
pub fn write_json_atomic_pretty<T, P>(path: P, value: &T) -> Result<()>
where
T: serde::Serialize,
P: AsRef<Path>,
{
write_json_atomic_inner(path.as_ref(), value, true)
}
fn write_json_atomic_inner<T>(path: &Path, value: &T, pretty: bool) -> Result<()>
where
T: serde::Serialize,
{
persist_atomic(path, |tmp| {
if pretty {
serde_json::to_writer_pretty(tmp, value).context("Failed to serialize JSON")
} else {
serde_json::to_writer(tmp, value).context("Failed to serialize JSON")
}
})
}
pub fn write_string_atomic<P: AsRef<Path>>(path: P, contents: &str) -> Result<()> {
persist_atomic(path.as_ref(), |tmp| {
tmp.write_all(contents.as_bytes())
.context("Failed to write file contents")
})
}
fn persist_atomic(
path: &Path,
write: impl FnOnce(&mut tempfile::NamedTempFile) -> Result<()>,
) -> Result<()> {
let dir = path.parent().unwrap_or_else(|| Path::new("."));
std::fs::create_dir_all(dir)
.with_context(|| format!("Failed to create directory: {}", dir.display()))?;
let mut tmp = tempfile::NamedTempFile::new_in(dir)
.with_context(|| format!("Failed to create temp file in: {}", dir.display()))?;
write(&mut tmp)?;
tmp.as_file().sync_all().ok();
tmp.persist(path)
.with_context(|| format!("Failed to persist file: {}", path.display()))?;
Ok(())
}
pub fn read_json<P: AsRef<Path>>(path: P) -> Result<Vec<Value>> {
let file = File::open(path.as_ref())
.with_context(|| format!("Failed to open file: {}", path.as_ref().display()))?;
let file_size = file.metadata().ok().map(|m| m.len() as usize).unwrap_or(0);
let mut contents = String::with_capacity(file_size);
let mut reader = BufReader::with_capacity(buffer::FILE_READ_BUFFER, file);
reader
.read_to_string(&mut contents)
.with_context(|| format!("Failed to read file: {}", path.as_ref().display()))?;
let obj: Value = serde_json::from_str(&contents).with_context(|| {
format!(
"Failed to parse JSON from file: {}",
path.as_ref().display()
)
})?;
Ok(vec![obj])
}
pub fn count_lines(text: &str) -> usize {
if text.is_empty() {
return 0;
}
let newline_count = bytecount::count(text.as_bytes(), b'\n');
if text.ends_with('\n') {
newline_count
} else {
newline_count + 1
}
}
pub fn save_json_pretty<P: AsRef<Path>>(path: P, value: &Value) -> Result<()> {
let json_str = serde_json::to_string_pretty(value).context("Failed to serialize JSON")?;
std::fs::write(path.as_ref(), json_str)
.with_context(|| format!("Failed to write file: {}", path.as_ref().display()))?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::fs::File;
use std::io::Write;
use tempfile::tempdir;
#[test]
fn test_count_lines_empty() {
assert_eq!(count_lines(""), 0);
}
#[test]
fn test_count_lines_single_line_no_newline() {
assert_eq!(count_lines("hello"), 1);
}
#[test]
fn test_count_lines_single_line_with_newline() {
assert_eq!(count_lines("hello\n"), 1);
}
#[test]
fn test_count_lines_multiple_lines() {
assert_eq!(count_lines("line1\nline2\nline3"), 3);
}
#[test]
fn test_count_lines_multiple_lines_with_newline() {
assert_eq!(count_lines("line1\nline2\nline3\n"), 3);
}
#[test]
fn test_count_lines_empty_lines() {
assert_eq!(count_lines("line1\n\nline3"), 3);
assert_eq!(count_lines("\n\n\n"), 3);
}
#[test]
fn test_read_jsonl_valid() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.jsonl");
let mut file = File::create(&file_path).unwrap();
writeln!(file, r#"{{"key1": "value1"}}"#).unwrap();
writeln!(file, r#"{{"key2": "value2"}}"#).unwrap();
writeln!(file, r#"{{"key3": "value3"}}"#).unwrap();
let result = read_jsonl(&file_path).unwrap();
assert_eq!(result.len(), 3);
assert_eq!(result[0]["key1"], "value1");
assert_eq!(result[1]["key2"], "value2");
assert_eq!(result[2]["key3"], "value3");
}
#[test]
fn test_read_jsonl_with_empty_lines() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.jsonl");
let mut file = File::create(&file_path).unwrap();
writeln!(file, r#"{{"key1": "value1"}}"#).unwrap();
writeln!(file).unwrap();
writeln!(file, r#"{{"key2": "value2"}}"#).unwrap();
writeln!(file, " ").unwrap();
writeln!(file, r#"{{"key3": "value3"}}"#).unwrap();
let result = read_jsonl(&file_path).unwrap();
assert_eq!(result.len(), 3);
}
#[test]
fn test_read_jsonl_empty_file() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("empty.jsonl");
File::create(&file_path).unwrap();
let result = read_jsonl(&file_path).unwrap();
assert_eq!(result.len(), 0);
}
#[test]
fn test_read_jsonl_invalid_json() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("invalid.jsonl");
let mut file = File::create(&file_path).unwrap();
writeln!(file, "not valid json").unwrap();
let result = read_jsonl(&file_path);
assert!(result.is_err());
}
#[test]
fn test_read_jsonl_nonexistent_file() {
let result = read_jsonl("/nonexistent/path/file.jsonl");
assert!(result.is_err());
}
#[test]
fn test_read_json_valid() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.json");
let mut file = File::create(&file_path).unwrap();
write!(file, r#"{{"key": "value", "number": 42}}"#).unwrap();
let result = read_json(&file_path).unwrap();
assert_eq!(result.len(), 1);
assert_eq!(result[0]["key"], "value");
assert_eq!(result[0]["number"], 42);
}
#[test]
fn test_read_json_array() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("array.json");
let mut file = File::create(&file_path).unwrap();
write!(file, r#"[1, 2, 3, 4, 5]"#).unwrap();
let result = read_json(&file_path).unwrap();
assert_eq!(result.len(), 1);
assert!(result[0].is_array());
assert_eq!(result[0].as_array().unwrap().len(), 5);
}
#[test]
fn test_read_json_invalid() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("invalid.json");
let mut file = File::create(&file_path).unwrap();
write!(file, "not valid json").unwrap();
let result = read_json(&file_path);
assert!(result.is_err());
}
#[test]
fn test_read_json_nonexistent_file() {
let result = read_json("/nonexistent/path/file.json");
assert!(result.is_err());
}
#[test]
fn test_count_lines_unicode() {
assert_eq!(count_lines("こんにちは"), 1);
assert_eq!(count_lines("line1 😀\nline2 🎉\nline3 🚀"), 3);
}
#[test]
fn test_count_lines_windows_line_endings() {
assert_eq!(count_lines("line1\r\nline2\r\nline3"), 3);
}
#[test]
fn test_read_jsonl_large_objects() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("large.jsonl");
let mut file = File::create(&file_path).unwrap();
let large_obj = json!({
"field1": "value1",
"field2": "value2",
"nested": {
"a": 1,
"b": 2,
"c": [1, 2, 3, 4, 5]
}
});
writeln!(file, "{}", serde_json::to_string(&large_obj).unwrap()).unwrap();
let result = read_jsonl(&file_path).unwrap();
assert_eq!(result.len(), 1);
assert_eq!(result[0]["field1"], "value1");
assert_eq!(result[0]["nested"]["c"].as_array().unwrap().len(), 5);
}
#[test]
fn test_read_json_nested_structure() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("nested.json");
let nested = json!({
"level1": {
"level2": {
"level3": {
"value": "deep"
}
}
}
});
let mut file = File::create(&file_path).unwrap();
write!(file, "{}", serde_json::to_string(&nested).unwrap()).unwrap();
let result = read_json(&file_path).unwrap();
assert_eq!(result.len(), 1);
assert_eq!(result[0]["level1"]["level2"]["level3"]["value"], "deep");
}
#[test]
fn test_count_lines_only_newlines() {
assert_eq!(count_lines("\n"), 1);
assert_eq!(count_lines("\n\n"), 2);
assert_eq!(count_lines("\n\n\n"), 3);
}
#[test]
fn test_count_lines_mixed_content() {
assert_eq!(count_lines("a\n \nc"), 3);
assert_eq!(count_lines(" hello \n world "), 2);
}
}