pub use crate::config::Config;
pub use crate::config::PrettyPrinterConfig;
use core::fmt::Write as _;
use std::fs;
use std::fs::File;
use std::io;
use std::io::Read;
use std::io::Seek;
use std::io::Write;
use std::path;
use beancount_pretty_printer::PrettyPrinter;
use beancount_types::Acc;
use beancount_types::Balance;
use beancount_types::Close;
use beancount_types::Commodity;
use beancount_types::Directive;
use beancount_types::Open;
use beancount_types::Price;
use camino::Utf8Path;
use camino::Utf8PathBuf;
use itertools::Itertools as _;
use miette::Diagnostic;
use rayon::prelude::IntoParallelIterator as _;
use rayon::prelude::ParallelIterator as _;
use rayon::slice::ParallelSliceMut as _;
use snafu::Backtrace;
use snafu::ResultExt as _;
use snafu::Snafu;
use time::Date;
use tracing::warn;
mod config;
#[derive(Debug, Diagnostic, Snafu)]
#[snafu(display("encountered multiple errors while writing tree of accounts"))]
pub struct Error {
#[related]
errors: Vec<ErrorKind>,
}
#[derive(Debug)]
pub struct TreeWriter {
config: Config,
}
impl TreeWriter {
pub fn new(config: Config) -> Self {
Self { config }
}
}
impl TreeWriter {
pub fn write_directives(&self, directives: Vec<Directive>) -> Result<()> {
let pretty_printer_config = match self.config.pretty_printer {
PrettyPrinterConfig::GloballyDerived => Some(
beancount_pretty_printer::Config::derive_from_directives(&directives),
),
PrettyPrinterConfig::LocallyDerived => None,
PrettyPrinterConfig::Static(config) => Some(config),
};
let accounts = directives
.into_iter()
.filter_map(|directive| {
Some((file_for(&self.config.output_path, &directive)?, directive))
})
.into_group_map();
let errors = accounts
.into_par_iter()
.map(|(file, directives)| self.write_month(pretty_printer_config, &file, directives))
.filter_map(Result::err)
.collect::<Vec<_>>();
if errors.is_empty() {
Ok(())
} else {
Err(Error { errors })
}
}
}
impl TreeWriter {
fn write_month(
&self,
pretty_printer_config: Option<beancount_pretty_printer::Config>,
path: &Utf8Path,
mut directives: Vec<Directive>,
) -> Result<(), ErrorKind> {
directives.par_sort_by(|left, right| {
left.timestamp().zip(right.timestamp()).map_or_else(
|| {
left.date()
.cmp(&right.date())
.then_with(|| left.kind().cmp(&right.kind()))
},
|(left, right)| left.cmp(&right),
)
});
let pretty_printer_config = pretty_printer_config.unwrap_or_else(|| {
beancount_pretty_printer::Config::derive_from_directives(&directives)
});
let mut buffer = Vec::new();
PrettyPrinter::unbuffered(pretty_printer_config, &mut buffer)
.print_directives(&directives)
.and_then(|_| update_file(path, &buffer))
.context(PrettyPrintingSnafu { file: path })?;
Ok(())
}
}
fn file_for(base: &Utf8Path, directive: &Directive) -> Option<Utf8PathBuf> {
match directive {
Directive::Balance(Balance { date, account, .. })
| Directive::Close(Close { date, account, .. })
| Directive::Open(Open { date, account, .. }) => Some(account_file(base, account, *date)),
Directive::Price(Price { quote, .. }) => Some(price_file(base, quote)),
Directive::Transaction(transaction) => {
if let Some(account) = transaction.main_account() {
Some(account_file(base, account, transaction.date))
} else {
warn!(
?transaction,
"ignoring transaction since it has no postings"
);
None
}
}
}
}
fn account_file(base: &Utf8Path, account: &Acc, date: Date) -> Utf8PathBuf {
let year = date.year();
let month = date.month();
let month = month as u8;
let separator = path::MAIN_SEPARATOR;
let mut path = base.to_string();
let additional = account.len() + 4 + 2 + 5 + 3; path.reserve(additional);
account
.segments()
.for_each(|segment| write!(path, "{separator}{segment}").unwrap());
write!(path, "{separator}{year}{separator}{month:02}.bean").unwrap();
Utf8PathBuf::from(path)
}
fn price_file(base: &Utf8Path, quote: &Commodity) -> Utf8PathBuf {
let mut path = base.join("prices");
path.push(&**quote);
path.set_extension("bean");
path
}
#[derive(Debug, Diagnostic, Snafu)]
enum ErrorKind {
#[snafu(display("encountered unsupported directive while indexing: {directive:?}"))]
IndexingDirective {
backtrace: Backtrace,
directive: Directive,
},
#[snafu(display("error while pretty printing directives to {file}"))]
PrettyPrinting {
backtrace: Backtrace,
file: Utf8PathBuf,
source: io::Error,
},
}
type Result<T, E = Error> = core::result::Result<T, E>;
fn ensure_directory(directory: &Utf8Path) -> io::Result<()> {
fs::create_dir_all(directory)
}
fn open_file(file_path: &Utf8Path) -> io::Result<File> {
ensure_directory(file_path.parent().unwrap())?;
File::options()
.create(true)
.read(true)
.truncate(true)
.write(true)
.open(file_path)
}
fn update_file(path: &Utf8Path, buffer: &[u8]) -> io::Result<()> {
let mut file = open_file(path)?;
let buffer_len = u64::try_from(buffer.len()).expect("should always work");
let file_size = file.metadata()?.len();
if buffer_len == file_size {
let mut current_contents = Vec::with_capacity(buffer.len());
file.read_to_end(&mut current_contents)?;
if current_contents == buffer {
tracing::debug!(%path, "skipping write of unchanged file");
return Ok(());
}
file.rewind()?;
}
file.write_all(buffer)?;
file.set_len(buffer_len)?;
file.sync_all()?;
tracing::debug!(%path, "successfully wrote account file");
Ok(())
}