From 228894e871369ca25514c3272a26ae98a62bb126 Mon Sep 17 00:00:00 2001 From: Erik T Date: Fri, 28 Jul 2023 15:05:07 -0400 Subject: [PATCH] added logic to select layer based on rarity --- Cargo.lock | 21 ++++++ Cargo.toml | 1 + src/cli/generate.rs | 175 ++++++++++++++++++++++++++------------------ src/cli/init.rs | 18 +++++ src/cli/mod.rs | 9 +++ src/cmd.rs | 3 + src/fs/dir.rs | 11 ++- src/main.rs | 1 + 8 files changed, 165 insertions(+), 74 deletions(-) create mode 100644 src/cli/init.rs diff --git a/Cargo.lock b/Cargo.lock index efd9c6b..599079d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -441,6 +441,7 @@ dependencies = [ "rand", "serde", "serde_json", + "thiserror", ] [[package]] @@ -847,6 +848,26 @@ version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1141d4d61095b28419e22cb0bbf02755f5e54e0526f97f1e3d1d160e60885fb" +[[package]] +name = "thiserror" +version = "1.0.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "611040a08a0439f8248d1990b111c95baa9c704c805fa1f62104b39655fd7f90" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "090198534930841fab3a5d1bb637cde49e339654e606195f8d9c76eeb081dc96" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.27", +] + [[package]] name = "threadpool" version = "1.8.1" diff --git a/Cargo.toml b/Cargo.toml index 8d0ee69..e0bae60 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,3 +16,4 @@ log = "0.4.19" rand = "0.8.5" serde = { version = "1.0.136", features = ["derive"] } serde_json = "1.0" +thiserror = "1.0.44" diff --git a/src/cli/generate.rs b/src/cli/generate.rs index 6f1ee28..32548bd 100644 --- a/src/cli/generate.rs +++ b/src/cli/generate.rs @@ -3,6 +3,7 @@ use console::style; use image::imageops::overlay; use rand::distributions::WeightedIndex; use rand::prelude::*; +use rand::seq::SliceRandom; use std::collections::HashMap; use std::collections::{BTreeMap, BTreeSet}; @@ -13,12 +14,16 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use crate::cli::utils::Cmd; +use crate::cli::ConversionError; use crate::constants::{ASSETS_INPUT, ASSETS_OUTPUT, METADATA_OUTPUT, PALETTE_EMOJI}; -use crate::fs::dir::Dir; +use crate::fs::dir::{Dir, DirError}; use crate::models::metadata::Metadata; +// TODO: Abstract a lot of the types away to structs and types + #[derive(Debug, Clone, Parser)] pub struct GenerateArgs { + /// Number of assets to generate #[clap(short, long)] count: u128, @@ -55,15 +60,7 @@ impl Cmd for GenerateArgs { eyre::bail!("Directory {:?} does not exist", cwd) } - let trait_layers = load_layers(input); - - // for s in subdirs { - // let dir = fs::read_dir(s)?; - // - // for item in dir { - // println!("{:?}", item.unwrap().path().display()) - // } - // } + let trait_layers = load_layers(input)?; // Create the assets output folder if it does not exist if !assets.exists() { @@ -75,11 +72,31 @@ impl Cmd for GenerateArgs { // fs::create_dir_all(&metadata)?; } - for i in 0..count { - todo!() - } + let trait_layer_keys: Vec = trait_layers.keys().cloned().collect(); + println!("{:?}", trait_layer_keys); + + for _ in 0..count { + let selected_layers: Vec<&Box> = trait_layer_keys + .iter() + .map(|trait_type| { + let mut rng = thread_rng(); + + let layer = match trait_layers.get(trait_type) { + Some(l) => l.choose_weighted(&mut rng, |x| x.rarity).unwrap(), + // TODO: Add a more descriptive error + None => eyre::bail!("Could not find layers for trait type"), + }; - Ok(()) + Ok(layer) + }) + .map(|l| l.unwrap()) + .collect(); + + println!("{:?}", selected_layers); + + let asset = create_artwork(&selected_layers); + } + todo!() } } @@ -90,81 +107,93 @@ struct Attribute { #[derive(Debug)] struct Layer { + /// Trait type of the layer (e.g., background, foreground, body, etc.) trait_type: String, - trait_name: String, + /// Value of the relative trait type + value: String, + /// Probability of being selected relative to other layers rarity: u32, } -// impl Layer { -// fn new(self, id: u32, trait_name: String, rarity: u32) -> Self { -// Layer { -// id, -// trait_name, -// rarity, -// } -// } -// } - -fn create_artwork(layers: &[&Arc]) { +impl Layer { + fn new(trait_type: String, value: String, rarity: u32) -> Self { + Layer { + trait_type, + value, + rarity, + } + } +} + +fn create_artwork(layers: &[&Box]) { todo!() } -fn encode_combination(layers: &[&Arc]) -> eyre::Result { +fn encode_combination(layers: &[&Box]) -> eyre::Result { todo!() } -type ArtLayers = Vec>>; +type TraitLayers = BTreeMap>>; -fn load_layers(input_dir: PathBuf) -> eyre::Result { +fn load_layers(input_dir: PathBuf) -> eyre::Result { let subdirs = Dir::read_dir(input_dir)?.contents; - // This flow actually might be better written imperatively so that error handling is easier - let trait_layers = subdirs - .into_iter() - .map(|subdir| { - let trait_type = subdir.file_stem().unwrap(); - - // Logic for removing the layer order prefix. Still figuring out the order of - // operations of this code. - // If this is included, the `trait_type` Layer field needs to be updated to - // `trait_type.to_owned().into_string().unwrap()` - // let trait_type = trait_type.to_owned().into_string().unwrap(); - // let [_, trait_type]: [&str; 2] = trait_type - // .split("_") - // .collect::>() - // .try_into() - // .unwrap(); - - let subdir = fs::read_dir(&subdir).unwrap(); - subdir - .into_iter() - .map(|file| { - // Get the file from within each subdirectory - let file = file.unwrap().file_name(); - // Create a Path from the file - let file_path = Path::new(&file); - // Get the stem from the Path and convert it to a String - let file_stem = file_path - .file_stem() - .unwrap() - .to_owned() - .into_string() - .unwrap(); - - // Create a new Arc and return - Box::new(Layer { - trait_type: trait_type.to_owned().into_string().unwrap(), - trait_name: file_stem, - rarity: 1, - }) - }) - .collect::>>() - }) - .collect::(); + let mut trait_layers: TraitLayers = BTreeMap::new(); + + for subdir in subdirs { + let trait_type = subdir + .file_stem() + .ok_or(DirError::FileStemError( + "Error reading file stem".to_string(), + ))? + .to_owned() + .into_string() + .map_err(|_| { + // TODO: Update the following error to be more descriptive + ConversionError::OsStringToStringError( + "Failed to convert OsString to String".to_string(), + ) + })?; + + let subdir = fs::read_dir(&subdir)?; + + let mut subdir_layers: Vec> = vec![]; + + for file in subdir { + let file = file?.file_name(); + + let file_path = Path::new(&file); + + let trait_value = file_path + .file_stem() + .ok_or(DirError::FileStemError( + "Error reading file stem.".to_string(), + ))? + .to_owned() + .into_string() + // TODO: Update the following error to be more descriptive + .map_err(|_| { + ConversionError::OsStringToStringError( + "Failed to convert OsString to String".to_string(), + ) + })?; + + let rarity = 1; + + // Cloning since I need trait_type later as well + let trait_type = trait_type.clone(); + + let layer = Box::new(Layer::new(trait_type, trait_value, rarity)); + + subdir_layers.push(layer); + } + + trait_layers.insert(trait_type, subdir_layers); + } Ok(trait_layers) } -fn generate_combinations(layers: &[&Arc]) -> eyre::Result<()> { +fn generate_combinations(layers: &[&Box]) -> eyre::Result<()> { todo!() } diff --git a/src/cli/init.rs b/src/cli/init.rs new file mode 100644 index 0000000..7e73f5f --- /dev/null +++ b/src/cli/init.rs @@ -0,0 +1,18 @@ +use clap::{Parser, ValueHint}; +use std::path::PathBuf; + +use crate::cli::utils::Cmd; + +// TODO: Add struct fields and run logic + +#[derive(Debug, Clone, Parser)] +pub struct InitArgs { + #[clap(short, long, value_hint = ValueHint::FilePath, value_name = "PATH", default_value = ".")] + root: PathBuf, +} + +impl Cmd for InitArgs { + fn run(self) -> eyre::Result<()> { + todo!() + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index b8cded5..8c0bff2 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -1,2 +1,11 @@ pub mod generate; +pub mod init; pub mod utils; + +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum ConversionError { + #[error("failed to convert OsString to String: {0}")] + OsStringToStringError(String), +} diff --git a/src/cmd.rs b/src/cmd.rs index 75c722c..146f108 100644 --- a/src/cmd.rs +++ b/src/cmd.rs @@ -1,4 +1,5 @@ use crate::cli::generate::GenerateArgs; +use crate::cli::init::InitArgs; use clap::{Parser, Subcommand}; #[derive(Debug, Parser)] @@ -11,6 +12,8 @@ pub struct Opts { #[derive(Debug, Subcommand)] #[clap(about = "Generate your latest profile picture.")] pub enum Subcommands { + /// Initialize new project + Init(InitArgs), /// Generate assets Generate(GenerateArgs), } diff --git a/src/fs/dir.rs b/src/fs/dir.rs index c8edc0b..5ba3e7b 100644 --- a/src/fs/dir.rs +++ b/src/fs/dir.rs @@ -1,6 +1,7 @@ use log::*; use std::fs; use std::path::PathBuf; +use thiserror::Error; pub struct Dir { pub contents: Vec, @@ -18,8 +19,16 @@ impl Dir { .collect::, _>>()?; // Sort the subdirectories in alphanumeric order - contents.sort(); + // This is unnecessary since we're using a BTreeMap but I'm keeping it in case we move to a + // Vec + // contents.sort(); Ok(Dir { contents, path }) } } + +#[derive(Error, Debug)] +pub enum DirError { + #[error("failed to get file stem: {0}")] + FileStemError(String), +} diff --git a/src/main.rs b/src/main.rs index f5b7f8f..c43a4f4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,7 @@ fn main() -> eyre::Result<()> { let opts = Opts::parse(); match opts.sub { + Subcommands::Init(cmd) => cmd.run(), Subcommands::Generate(cmd) => cmd.run(), } }