From 19b484a7445f69642ef575a8847f98235d3ea475 Mon Sep 17 00:00:00 2001
From: sxyazi <sxyazi@gmail.com>
Date: Mon, 28 Apr 2025 21:41:28 +0800
Subject: [PATCH] feat: new `fs.calc_size()` API

---
 yazi-fs/src/calculator.rs             | 11 ++++++++---
 yazi-plugin/src/bindings/mod.rs       |  2 +-
 yazi-plugin/src/bindings/receiver.rs  |  0
 yazi-plugin/src/fs/fs.rs              | 26 ++++++++++++++++++++++++--
 yazi-scheduler/src/prework/prework.rs |  2 +-
 5 files changed, 34 insertions(+), 7 deletions(-)
 create mode 100644 yazi-plugin/src/bindings/receiver.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/mod.rs b/yazi-plugin/src/bindings/mod.rs
index 7a8b9765..c2a77b4c 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!(cha chan icon image input layer mouse permit range receiver);
diff --git a/yazi-plugin/src/bindings/receiver.rs b/yazi-plugin/src/bindings/receiver.rs
new file mode 100644
index 00000000..e69de29b
diff --git a/yazi-plugin/src/fs/fs.rs b/yazi-plugin/src/fs/fs.rs
index f26db46f..37355859 100644
--- a/yazi-plugin/src/fs/fs.rs
+++ b/yazi-plugin/src/fs/fs.rs
@@ -1,8 +1,8 @@
 use globset::GlobBuilder;
-use mlua::{ExternalError, ExternalResult, Function, IntoLua, IntoLuaMulti, Lua, Table, Value};
+use mlua::{ExternalError, ExternalResult, Function, IntoLua, IntoLuaMulti, Lua, Table, UserData, UserDataMethods, Value};
 use tokio::fs;
 use yazi_binding::{Error, Url, UrlRef};
-use yazi_fs::{mounts::PARTITIONS, remove_dir_clean};
+use yazi_fs::{SizeCalculator, mounts::PARTITIONS, remove_dir_clean};
 
 use crate::{Composer, bindings::Cha, file::File};
 
@@ -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,27 @@ 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 {
+		struct Wrapper(SizeCalculator);
+		impl UserData for Wrapper {
+			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),
+					}
+				});
+			}
+		}
+
+		match SizeCalculator::new(&*url).await {
+			Ok(it) => Wrapper(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();