diff options
author | Soispha <soispha@vhack.eu> | 2024-01-05 13:26:31 +0100 |
---|---|---|
committer | Soispha <soispha@vhack.eu> | 2024-01-05 13:26:31 +0100 |
commit | d33998b3ae3da72240d1dcc9f7766d42b94e0999 (patch) | |
tree | 69795f58c23d88fccebf97e1519a18ebc253520e /sys/nixpkgs/pkgs/ytc/src | |
parent | fix(hm/pkgs/scr/ytcc/filter_comments.sh): Rename export script to 'comments' (diff) | |
download | nixos-config-d33998b3ae3da72240d1dcc9f7766d42b94e0999.tar.gz nixos-config-d33998b3ae3da72240d1dcc9f7766d42b94e0999.zip |
feat(hm/pkgs/scr/ytcc/ytc): Rewrite in Rust
Diffstat (limited to 'sys/nixpkgs/pkgs/ytc/src')
-rw-r--r-- | sys/nixpkgs/pkgs/ytc/src/downloader.rs | 146 | ||||
-rw-r--r-- | sys/nixpkgs/pkgs/ytc/src/main.rs | 188 |
2 files changed, 334 insertions, 0 deletions
diff --git a/sys/nixpkgs/pkgs/ytc/src/downloader.rs b/sys/nixpkgs/pkgs/ytc/src/downloader.rs new file mode 100644 index 00000000..9509278c --- /dev/null +++ b/sys/nixpkgs/pkgs/ytc/src/downloader.rs @@ -0,0 +1,146 @@ +use std::{ + io::{stderr, stdout, Read}, + mem, + path::PathBuf, + process::Command, + sync::mpsc::{self, Receiver, Sender}, + thread::{self, JoinHandle}, +}; + +use anyhow::{bail, Context, Result}; +use log::debug; + +use crate::PlayThing; + +const YT_DLP_FLAGS: [&str; 12] = [ + "--format", + "bestvideo[height<=?1080]+bestaudio/best", + "--embed-chapters", + "--progress", + "--write-comments", + "--extractor-args", + "youtube:max_comments=150,all,100;comment_sort=top", + "--write-info-json", + "--sponsorblock-mark", + "default", + "--sponsorblock-remove", + "sponsor", +]; + +const CONCURRENT: u32 = 5; + +const DOWNLOAD_DIR: &str = "/tmp/ytcc"; + +pub struct Downloader { + sent: usize, + download_thread: JoinHandle<Result<()>>, + orx: Receiver<(PathBuf, Option<u32>)>, + itx: Option<Sender<PlayThing>>, + playspec: Vec<PlayThing>, +} + +impl Downloader { + pub fn new(mut playspec: Vec<PlayThing>) -> anyhow::Result<Downloader> { + let (itx, irx): (Sender<PlayThing>, Receiver<PlayThing>) = mpsc::channel(); + let (otx, orx) = mpsc::channel(); + let jh = thread::spawn(move || -> Result<()> { + while let Some(pt) = irx.recv().ok() { + debug!("Got '{}|{}' to be downloaded", pt.url, pt.id.unwrap_or(0)); + let path = download_url(&pt.url) + .with_context(|| format!("Failed to download url: '{}'", &pt.url))?; + otx.send((path, pt.id)).expect("Should not be dropped"); + } + debug!("Finished Downloading everything"); + Ok(()) + }); + + playspec.reverse(); + let mut output = Downloader { + sent: 0, + download_thread: jh, + orx, + itx: Some(itx), + playspec, + }; + if output.playspec.len() <= CONCURRENT as usize { + output.add(output.playspec.len() as u32)?; + } else { + output.add(CONCURRENT)?; + } + Ok(output) + } + + pub fn add(&mut self, number_to_add: u32) -> Result<()> { + debug!("Adding {} to be downloaded concurrently", number_to_add); + for _ in 0..number_to_add { + let pt = self.playspec.pop().context("No more playthings to pop")?; + self.itx.as_ref().expect("Should still be valid").send(pt)?; + } + Ok(()) + } + + /// Return the next video already downloaded, will block until the download is complete + pub fn next(&mut self) -> Option<(PathBuf, Option<u32>)> { + debug!("Requesting next output"); + match self.orx.recv() { + Ok(ok) => { + debug!("Output downloaded to: {}", ok.0.display()); + self.sent += 1; + if self.sent < self.playspec.len() { + debug!("Will add 1"); + self.add(1).ok()?; + } else { + debug!("Will drop sender"); + let itx = mem::take(&mut self.itx); + drop(itx) + } + debug!("Returning: {:#?}", ok); + Some(ok) + } + Err(err) => { + debug!("Recieved error while listening: {}", err); + None + } + } + } + pub fn drop(self) -> anyhow::Result<()> { + match self.download_thread.join() { + Ok(ok) => ok, + Err(err) => panic!("Can't join thread: '{:#?}'", err), + } + } +} + +fn download_url(url: &str) -> Result<PathBuf> { + let output_file = tempfile::NamedTempFile::new().context("Failed to create tempfile")?; + output_file + .as_file() + .set_len(0) + .context("Failed to truncate temp-file")?; + let mut yt_dlp = Command::new("yt-dlp"); + yt_dlp.current_dir(DOWNLOAD_DIR); + yt_dlp.stdout(stdout()); + yt_dlp.stderr(stderr()); + yt_dlp.args(YT_DLP_FLAGS); + yt_dlp.args([ + "--output", + "%(channel)s/%(title)s.%(ext)s", + url, + "--print-to-file", + "after_move:filepath", + ]); + yt_dlp.arg(output_file.path().as_os_str()); + let status = yt_dlp.status().context("Failed to run yt_dlp")?; + if let Some(code) = status.code() { + if code != 0 { + bail!("yt_dlp execution failed with error: '{}'", status); + } + } + let mut path = String::new(); + output_file + .as_file() + .read_to_string(&mut path) + .context("Failed to read output file temp file")?; + let path = path.trim(); + Ok(path.into()) +} diff --git a/sys/nixpkgs/pkgs/ytc/src/main.rs b/sys/nixpkgs/pkgs/ytc/src/main.rs new file mode 100644 index 00000000..5c7849b8 --- /dev/null +++ b/sys/nixpkgs/pkgs/ytc/src/main.rs @@ -0,0 +1,188 @@ +use std::{ + env, + fs::{self, canonicalize}, + io::{stderr, stdout}, + os::unix::fs::symlink, + path::PathBuf, + process::Command as StdCmd, +}; + +use anyhow::{bail, Context, Result}; +use clap::{Parser, Subcommand}; +use downloader::Downloader; +use log::debug; +use serde::Deserialize; + +mod downloader; + +const STATUS_PATH: &str = "ytcc/running"; + +/// A helper for downloading and playing youtube videos +#[derive(Parser, Debug)] +#[clap(author, version, about, long_about = None)] +pub struct Args { + #[command(subcommand)] + /// The subcommand to execute + pub subcommand: Command, +} +#[derive(Subcommand, Debug)] +pub enum Command { + #[clap(value_parser)] + /// Work based of ytcc ids + Id { + #[clap(value_parser)] + /// A list of ids to play + ids: Vec<u32>, + }, + #[clap(value_parser)] + /// Work based of raw youtube urls + Url { + #[clap(value_parser)] + /// A list of urls to play + urls: Vec<String>, + }, +} + +struct PlayThing { + url: String, + id: Option<u32>, +} + +#[derive(Deserialize)] +struct YtccListData { + url: String, + #[allow(unused)] + title: String, + #[allow(unused)] + description: String, + #[allow(unused)] + publish_date: String, + #[allow(unused)] + watch_date: Option<String>, + #[allow(unused)] + duration: String, + #[allow(unused)] + thumbnail_url: String, + #[allow(unused)] + extractor_hash: String, + id: u32, + #[allow(unused)] + playlists: Vec<YtccPlaylistData>, +} +#[derive(Deserialize)] +struct YtccPlaylistData { + #[allow(unused)] + name: String, + #[allow(unused)] + url: String, + #[allow(unused)] + reverse: bool, +} + +fn main() -> Result<()> { + let args = Args::parse(); + cli_log::init_cli_log!(); + + let playspec: Vec<PlayThing> = match args.subcommand { + Command::Id { ids } => { + let mut output = Vec::with_capacity(ids.len()); + for id in ids { + debug!("Adding {}", id); + let mut ytcc = StdCmd::new("ytcc"); + ytcc.args([ + "--output", + "json", + "list", + "--attributes", + "url", + "--ids", + id.to_string().as_str(), + ]); + let json = serde_json::from_slice::<Vec<YtccListData>>( + &ytcc.output().context("Failed to get url from id")?.stdout, + ) + .context("Failed to deserialize json output")?; + + if json.len() == 0 { + bail!("Could not find a video with id: {}", id); + } + assert_eq!(json.len(), 1); + let json = json.first().expect("Has only one element"); + + debug!("Id resolved to: '{}'", &json.url); + + output.push(PlayThing { + url: json.url.clone(), + id: Some(json.id), + }) + } + output + } + Command::Url { urls } => urls + .into_iter() + .map(|url| PlayThing { url, id: None }) + .collect(), + }; + + debug!("Initializing downloader"); + let mut downloader = Downloader::new(playspec)?; + + while let Some((path, id)) = downloader.next() { + debug!("Next path to play is: '{}'", path.display()); + let mut info_json = canonicalize(&path).context("Failed to canoncialize path")?; + info_json.set_extension("info.json"); + + if status_path()?.is_symlink() { + fs::remove_file(status_path()?).context("Failed to delete old status file")?; + } else { + bail!( + "The status path ('{}') is not a symlink!", + status_path()?.display() + ); + } + + symlink(info_json, status_path()?).context("Failed to symlink")?; + + let mut mpv = StdCmd::new("mpv"); + mpv.stdout(stdout()); + mpv.stderr(stderr()); + mpv.args(["--speed=2.7", "--volume=75"]); + mpv.arg(&path); + + let status = mpv.status().context("Failed to run mpv")?; + if let Some(code) = status.code() { + if let Some(id) = id { + if code == 0 { + println!("\x1b[32;1mMarking {} as watched!\x1b[0m", id); + fs::remove_file(&path)?; + let mut ytcc = StdCmd::new("ytcc"); + ytcc.stdout(stdout()); + ytcc.stderr(stderr()); + ytcc.args(["mark"]); + ytcc.arg(id.to_string()); + let status = ytcc.status().context("Failed to run ytcc")?; + if let Some(code) = status.code() { + if code != 0 { + bail!("Ytcc failed with status: {}", code); + } + } + } + } + debug!("Mpv exited with: {}", code); + } + } + downloader.drop()?; + + Ok(()) +} + +fn status_path() -> Result<PathBuf> { + let out: PathBuf = format!( + "{}/{}", + env::var("XDG_RUNTIME_DIR").expect("This should always exist"), + STATUS_PATH + ) + .into(); + fs::create_dir_all(&out.parent().expect("Parent should exist"))?; + Ok(out) +} |