From 2c99075ffa1efb83cc25ca1ddd3776d68671250f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=89=E5=92=B2=E9=9B=85=20=C2=B7=20Misaki=20Masa?=
 <sxyazi@gmail.com>
Date: Mon, 28 Apr 2025 21:56:08 +0800
Subject: [PATCH] feat: new `fs.calc_size()` API (#2691)

---
 yazi-fs/src/calculator.rs              | 11 ++++++++---
 yazi-plugin/src/bindings/calculator.rs | 15 +++++++++++++++
 yazi-plugin/src/bindings/mod.rs        |  2 +-
 yazi-plugin/src/fs/fs.rs               | 12 +++++++++++-
 yazi-scheduler/src/prework/prework.rs  |  2 +-
 5 files changed, 36 insertions(+), 6 deletions(-)
 create mode 100644 yazi-plugin/src/bindings/calculator.rs

diff --git a/yazi-fs/src/calculator.rs b/yazi-fs/src/calculator.rs
index 91ecb8ab..554274b5 100644
--- a/yazi-fs/src/calculator.rs
+++ b/yazi-fs/src/calculator.rs
@@ -13,7 +13,12 @@ pub enum SizeCalculator {
 impl SizeCalculator {
 	pub async fn new(path: impl AsRef<Path>) -> std::io::Result<Self> {
 		let p = path.as_ref().to_owned();
-		tokio::task::spawn_blocking(|| {
+		tokio::task::spawn_blocking(move || {
+			let meta = std::fs::symlink_metadata(&p)?;
+			if !meta.is_dir() {
+				return Ok(Self::Idle((VecDeque::new(), Some(meta.len()))));
+			}
+
 			let mut buf = VecDeque::from([Either::Right(std::fs::read_dir(p)?)]);
 			let size = Self::next_chunk(&mut buf);
 			Ok(Self::Idle((buf, size)))
@@ -87,8 +92,8 @@ impl SizeCalculator {
 			let Ok(ft) = ent.file_type() else { continue };
 			if ft.is_dir() {
 				buf.push_back(Either::Left(ent.path()));
-			} else {
-				size += ent.metadata().map(|m| m.len()).unwrap_or(0);
+			} else if let Ok(meta) = ent.metadata() {
+				size += meta.len();
 			}
 		}
 		Some(size)
diff --git a/yazi-plugin/src/bindings/calculator.rs b/yazi-plugin/src/bindings/calculator.rs
new file mode 100644
index 00000000..06fa002c
--- /dev/null
+++ b/yazi-plugin/src/bindings/calculator.rs
@@ -0,0 +1,15 @@
+use mlua::{IntoLuaMulti, UserData, UserDataMethods, Value};
+use yazi_binding::Error;
+
+pub struct SizeCalculator(pub yazi_fs::SizeCalculator);
+
+impl UserData for SizeCalculator {
+	fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
+		methods.add_async_method_mut("recv", |lua, mut me, ()| async move {
+			match me.0.next().await {
+				Ok(value) => (value, Value::Nil).into_lua_multi(&lua),
+				Err(e) => (Value::Nil, Error::Io(e)).into_lua_multi(&lua),
+			}
+		});
+	}
+}
diff --git a/yazi-plugin/src/bindings/mod.rs b/yazi-plugin/src/bindings/mod.rs
index 7a8b9765..8481d045 100644
--- a/yazi-plugin/src/bindings/mod.rs
+++ b/yazi-plugin/src/bindings/mod.rs
@@ -1,3 +1,3 @@
 #![allow(clippy::module_inception)]
 
-yazi_macro::mod_flat!(cha chan icon image input layer mouse permit range);
+yazi_macro::mod_flat!(calculator cha chan icon image input layer mouse permit range);
diff --git a/yazi-plugin/src/fs/fs.rs b/yazi-plugin/src/fs/fs.rs
index f26db46f..fd90d672 100644
--- a/yazi-plugin/src/fs/fs.rs
+++ b/yazi-plugin/src/fs/fs.rs
@@ -4,7 +4,7 @@ use tokio::fs;
 use yazi_binding::{Error, Url, UrlRef};
 use yazi_fs::{mounts::PARTITIONS, remove_dir_clean};
 
-use crate::{Composer, bindings::Cha, file::File};
+use crate::{Composer, bindings::{Cha, SizeCalculator}, file::File};
 
 pub fn compose(lua: &Lua) -> mlua::Result<Value> {
 	Composer::make(lua, 10, |lua, key| {
@@ -16,6 +16,7 @@ pub fn compose(lua: &Lua) -> mlua::Result<Value> {
 			b"create" => create(lua)?,
 			b"remove" => remove(lua)?,
 			b"read_dir" => read_dir(lua)?,
+			b"calc_size" => calc_size(lua)?,
 			b"expand_url" => expand_url(lua)?,
 			b"unique_name" => unique_name(lua)?,
 			b"partitions" => partitions(lua)?,
@@ -152,6 +153,15 @@ fn read_dir(lua: &Lua) -> mlua::Result<Function> {
 	})
 }
 
+fn calc_size(lua: &Lua) -> mlua::Result<Function> {
+	lua.create_async_function(|lua, url: UrlRef| async move {
+		match yazi_fs::SizeCalculator::new(&*url).await {
+			Ok(it) => SizeCalculator(it).into_lua_multi(&lua),
+			Err(e) => (Value::Nil, Error::Io(e)).into_lua_multi(&lua),
+		}
+	})
+}
+
 fn expand_url(lua: &Lua) -> mlua::Result<Function> {
 	lua.create_function(|_, value: Value| {
 		use yazi_fs::expand_path;
diff --git a/yazi-scheduler/src/prework/prework.rs b/yazi-scheduler/src/prework/prework.rs
index 9a978371..0e423b3a 100644
--- a/yazi-scheduler/src/prework/prework.rs
+++ b/yazi-scheduler/src/prework/prework.rs
@@ -73,7 +73,7 @@ impl Prework {
 				self.prog.send(TaskProg::Adv(task.id, 1, 0))?;
 			}
 			PreworkOp::Size(task) => {
-				let length = SizeCalculator::total(&task.target).await?;
+				let length = SizeCalculator::total(&task.target).await.unwrap_or(0);
 				task.throttle.done((task.target, length), |buf| {
 					{
 						let mut loading = self.size_loading.write();