v1.2.0: can now weight the random selection of shows

Cargo.lock

@@ -17,6 +17,15 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"
 
+[[package]]
+name = "android_system_properties"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "anyhow"
 version = "1.0.71"
@@ -116,8 +125,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ec837a71355b28f6556dbd569b37b3f363091c0bd4b2e735674521b4c5fd9bc5"
 dependencies = [
  "android-tzdata",
+ "iana-time-zone",
  "num-traits",
  "serde",
+ "winapi",
 ]
 
 [[package]]
@@ -193,6 +204,41 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "darling"
+version = "0.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0558d22a7b463ed0241e993f76f09f30b126687447751a8638587b864e4b3944"
+dependencies = [
+ "darling_core",
+ "darling_macro",
+]
+
+[[package]]
+name = "darling_core"
+version = "0.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab8bfa2e259f8ee1ce5e97824a3c55ec4404a0d772ca7fa96bf19f0752a046eb"
+dependencies = [
+ "fnv",
+ "ident_case",
+ "proc-macro2",
+ "quote",
+ "strsim",
+ "syn 2.0.18",
+]
+
+[[package]]
+name = "darling_macro"
+version = "0.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29a358ff9f12ec09c3e61fef9b5a9902623a695a46a917b07f269bff1445611a"
+dependencies = [
+ "darling_core",
+ "quote",
+ "syn 2.0.18",
+]
+
 [[package]]
 name = "derive_more"
 version = "0.99.17"
@@ -506,6 +552,12 @@ version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286"
 
+[[package]]
+name = "hex"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
+
 [[package]]
 name = "hmac"
 version = "0.12.1"
@@ -592,6 +644,35 @@ dependencies = [
  "tokio-native-tls",
 ]
 
+[[package]]
+name = "iana-time-zone"
+version = "0.1.57"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2fad5b825842d2b38bd206f3e81d6957625fd7f0a361e345c30e01a0ae2dd613"
+dependencies = [
+ "android_system_properties",
+ "core-foundation-sys",
+ "iana-time-zone-haiku",
+ "js-sys",
+ "wasm-bindgen",
+ "windows",
+]
+
+[[package]]
+name = "iana-time-zone-haiku"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "ident_case"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
+
 [[package]]
 name = "idna"
 version = "0.2.3"
@@ -637,6 +718,7 @@ checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
 dependencies = [
  "autocfg",
  "hashbrown 0.12.3",
+ "serde",
 ]
 
 [[package]]
@@ -1137,7 +1219,7 @@ checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
 
 [[package]]
 name = "screencap-bot"
-version = "1.1.0"
+version = "1.2.0"
 dependencies = [
  "anyhow",
  "dotenvy",
@@ -1150,6 +1232,7 @@ dependencies = [
  "rand",
  "regex",
  "serde",
+ "serde_with",
  "serde_yaml",
  "tempfile",
  "tokio",
@@ -1221,6 +1304,34 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "serde_with"
+version = "3.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f02d8aa6e3c385bf084924f660ce2a3a6bd333ba55b35e8590b321f35d88513"
+dependencies = [
+ "base64 0.21.2",
+ "chrono",
+ "hex",
+ "indexmap 1.9.3",
+ "serde",
+ "serde_json",
+ "serde_with_macros",
+ "time",
+]
+
+[[package]]
+name = "serde_with_macros"
+version = "3.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "edc7d5d3932fb12ce722ee5e64dd38c504efba37567f0c402f6ca728c3b8b070"
+dependencies = [
+ "darling",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.18",
+]
+
 [[package]]
 name = "serde_yaml"
 version = "0.9.22"
@@ -1285,6 +1396,12 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "strsim"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623"
+
 [[package]]
 name = "subtle"
 version = "2.5.0"
@@ -1705,6 +1822,15 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
+[[package]]
+name = "windows"
+version = "0.48.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f"
+dependencies = [
+ "windows-targets",
+]
+
 [[package]]
 name = "windows-sys"
 version = "0.42.0"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "screencap-bot"
-version = "1.1.0"
+version = "1.2.0"
 edition = "2021"
 authors = ["xenofem <xenofem@xeno.science>"]
 license = "MIT"
@@ -17,6 +17,7 @@ log = "0.4.19"
 rand = "0.8"
 regex = "1.8.4"
 serde = "1"
+serde_with = "3"
 serde_yaml = "0.9.22"
 tempfile = "3.6.0"
 tokio = { version = "1.28.2", features = ["full"] }

src/main.rs

@@ -4,7 +4,11 @@ use anyhow::{anyhow, Context};
 use config::Config;
 use lazy_static::lazy_static;
 use log::{debug, error, info};
-use rand::{distributions::Standard, seq::IteratorRandom, Rng};
+use rand::{
+    distributions::{Distribution, Standard, WeightedIndex},
+    seq::IteratorRandom,
+    Rng,
+};
 use shows::Shows;
 
 mod config;
@@ -36,13 +40,16 @@ async fn main() -> anyhow::Result<()> {
         return Err(anyhow!("Shows file is empty!"));
     }
 
+    let dist = WeightedIndex::new(shows.iter().map(|s| s.weight))
+        .context("Failed to load show weights")?;
+
     info!("Logging into cohost as {}", conf.cohost_email);
     let session = eggbug::Session::login(&conf.cohost_email, &conf.cohost_password)
         .await
         .context("Failed to login to cohost")?;
 
     loop {
-        let result = post_random_screencap(&conf, &shows, &session, &mut rng)
+        let result = post_random_screencap(&conf, &shows, &session, &dist, &mut rng)
             .await
             .context("Failed to post a random screencap");
 
@@ -65,19 +72,20 @@ async fn post_random_screencap<R: Rng>(
     conf: &Config,
     shows: &Shows,
     session: &eggbug::Session,
+    dist: &WeightedIndex<f32>,
     rng: &mut R,
 ) -> anyhow::Result<()> {
-    let (title, show) = shows.iter().choose(rng).unwrap();
+    let show = &shows[dist.sample(rng)];
     let episodes = show.episodes().with_context(|| {
         format!(
             "Failed to get episode list for show {} with path {}",
-            title,
+            show.title,
             show.path.display()
         )
     })?;
     let (num, file) = episodes.iter().choose(rng).unwrap();
 
-    let descriptor = shows::display_show_episode(title, show, *num);
+    let descriptor = shows::display_show_episode(show, *num);
 
     info!("Selected: {} - {}", descriptor, file.display());
 

src/shows/mod.rs

@@ -7,19 +7,32 @@ use std::{
 use anyhow::{anyhow, Context};
 use log::{debug, error};
 use serde::Deserialize;
+use serde_with::{serde_as, KeyValueMap};
 
 mod enumeration;
 
 #[derive(Deserialize)]
 pub struct Show {
+    #[serde(rename = "$key$")]
+    pub title: String,
     pub path: PathBuf,
     #[serde(default)]
     pub tags: Vec<String>,
     #[serde(default)]
     pub parts: HashMap<u32, String>,
+    #[serde(default = "default_weight")]
+    pub weight: f32,
 }
 
-pub type Shows = HashMap<String, Show>;
+fn default_weight() -> f32 {
+    1.0
+}
+
+pub type Shows = Vec<Show>;
+
+#[serde_as]
+#[derive(Deserialize)]
+struct ShowsWrapper(#[serde_as(as = "KeyValueMap<_>")] Shows);
 
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 pub enum EpisodeNumber {
@@ -31,17 +44,20 @@ pub enum EpisodeNumber {
 type Episodes = HashMap<EpisodeNumber, PathBuf>;
 
 pub fn load<P: AsRef<Path>>(shows_file: P) -> anyhow::Result<Shows> {
-    serde_yaml::from_reader(fs::File::open(shows_file).context("Failed to open shows file")?)
+    Ok(serde_yaml::from_reader::<_, ShowsWrapper>(
-        .context("Failed to parse YAML from shows file")
+        fs::File::open(shows_file).context("Failed to open shows file")?,
+    )
+    .context("Failed to parse YAML from shows file")?
+    .0)
 }
 
-pub fn display_show_episode(title: &str, show: &Show, episode: EpisodeNumber) -> String {
+pub fn display_show_episode(show: &Show, episode: EpisodeNumber) -> String {
     match episode {
-        EpisodeNumber::Standalone => title.to_string(),
+        EpisodeNumber::Standalone => show.title.to_string(),
-        EpisodeNumber::SingleSeason(n) => format!("{} episode {}", title, n),
+        EpisodeNumber::SingleSeason(n) => format!("{} episode {}", show.title, n),
         EpisodeNumber::MultiSeason(season, ep) => format!(
             "{} {} episode {}",
-            title,
+            show.title,
             show.parts
                 .get(&season)
                 .unwrap_or(&format!("season {}", season)),