From 709066c8c6a8ee1cc1f4dfbc032cfaad141e12f8 Mon Sep 17 00:00:00 2001
From: Benedikt Peetz <benedikt.peetz@b-peetz.de>
Date: Sat, 19 Oct 2024 14:28:32 +0200
Subject: fix(yt_dlp/lib/progress_hook): Estimate `total_byte_size` better

This still is sort of weird, because the total byte size changes whilst
downloading, but it is still immensely better than just putting a `0`
there.
---
 crates/yt_dlp/src/lib.rs | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

(limited to 'crates')

diff --git a/crates/yt_dlp/src/lib.rs b/crates/yt_dlp/src/lib.rs
index 4e35cb0..9d1a0c5 100644
--- a/crates/yt_dlp/src/lib.rs
+++ b/crates/yt_dlp/src/lib.rs
@@ -20,7 +20,7 @@ use std::{path::PathBuf, sync::Once};
 use crate::{duration::Duration, logging::setup_logging, wrapper::info_json::InfoJson};
 
 use bytes::Bytes;
-use log::{info, log_enabled, warn, Level};
+use log::{info, log_enabled, Level};
 use pyo3::types::{PyString, PyTuple, PyTupleMethods};
 use pyo3::{
     pyfunction,
@@ -102,6 +102,7 @@ signal.signal(signal.SIGINT, signal.SIG_DFL)",
 pub fn progress_hook(py: Python<'_>, input: &Bound<'_, PyDict>) -> PyResult<()> {
     // Only add the handler, if the log-level is higher than Debug (this avoids covering debug
     // messages).
+    // FIXME: We should actually just find a way to not cover printed messages. <2024-10-19>
     if log_enabled!(Level::Debug) {
         return Ok(());
     }
@@ -219,21 +220,19 @@ pub fn progress_hook(py: Python<'_>, input: &Bound<'_, PyDict>) -> PyResult<()>
             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
+                    let maybe_estimate = default_get!(as_u64, 0, "total_bytes_estimate");
+
+                    if maybe_estimate == 0 {
+                        // The download speed should be in bytes per second and the eta in seconds.
+                        // Thus multiplying them gets us the raw bytes (which were estimated by `yt_dlp`, from their `info.json`)
+                        let bytes_still_needed = (speed * eta).ceil() as u64;
+
+                        downloaded_bytes + bytes_still_needed
+                    } else {
+                        maybe_estimate
+                    }
                 } else {
                     total_bytes
-                    // FIXME: The solution below _would_ work, if we had a info_json with the
-                    // values. <2024-08-24>
-                    //
-                    // let duration = default_get! {as_f64, 0.0, "duration"}.ceil() as u64;
-                    // // TODO: yt_dlp gets this from the format
-                    // let tbr = default_get! {as_f64, 0.0, "tbr"}.ceil() as u64;
-                    // duration * tbr * (1000 / 8)
                 }
             };
             let percent: f64 = {
-- 
cgit 1.4.1