chore(yt_dlp/progress_hook): Also consider the `total_bytes_estimate` field

1 files changed, 14 insertions, 3 deletions

diff --git a/yt_dlp/src/lib.rs b/yt_dlp/src/lib.rs
index 27f1a58..e0cb590 100644
--- a/yt_dlp/src/lib.rs
+++ b/yt_dlp/src/lib.rs
@@ -14,7 +14,7 @@ use std::{path::PathBuf, sync::Once};
 
 use crate::{duration::Duration, logging::setup_logging, wrapper::info_json::InfoJson};
 
-use log::info;
+use log::{info, warn};
 use pyo3::types::{PyString, PyTuple, PyTupleMethods};
 use pyo3::{
     pyfunction,
@@ -200,8 +200,19 @@ pub fn progress_hook<'a>(py: Python, input: Bound<'_, PyDict>) -> PyResult<()> {
             let speed = default_get! {as_f64, 0.0, "speed"};
 
             let downloaded_bytes = get! {is_u64, as_u64, "downloaded_bytes"};
-            let total_bytes = default_get!(as_u64, 0, "total_bytes");
-
+            let total_bytes = {
+                let total_bytes = default_get!(as_u64, 0, "total_bytes");
+                if total_bytes == 0 {
+                    let estimate = default_get!(as_u64, 0, "total_bytes_estimate");
+                    warn!(
+                        "The video does not have a total_byte count, using an estimate of '{}'",
+                        estimate
+                    );
+                    estimate
+                } else {
+                    total_bytes
+                }
+            };
             let percent: f64 = {
                 if total_bytes == 0 {
                     100.0