Skip to content

Commit

Permalink
Merge pull request #381 from ReagentX/feat/cs/support-genmoji-prompts
Browse files Browse the repository at this point in the history 
Feat/cs/support genmoji prompts
  • Loading branch information
@ReagentX
ReagentX authored Nov 23, 2024
2 parents 1ccf091 + 0054080 commit d1c42af
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 23 deletions.
4 changes: 4 additions & 0 deletions imessage-database/src/tables/attachment.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ pub struct Attachment {
/// `true` if the attachment was a sticker, else `false`
pub is_sticker: bool,
pub hide_attachment: i32,
/// The prompt used to generate a Genmoji
pub emoji_description: Option<String>,
/// Auxiliary data to denote that an attachment has been copied
pub copied_path: Option<PathBuf>,
}
Expand All @@ -90,6 +92,7 @@ impl Table for Attachment {
total_bytes: row.get("total_bytes").unwrap_or_default(),
is_sticker: row.get("is_sticker").unwrap_or(false),
hide_attachment: row.get("hide_attachment").unwrap_or(0),
emoji_description: row.get("emoji_image_short_description").unwrap_or(None),
copied_path: None,
})
}
Expand Down Expand Up @@ -468,6 +471,7 @@ mod tests {
total_bytes: 100,
is_sticker: false,
hide_attachment: 0,
emoji_description: None,
copied_path: None,
}
}
Expand Down
23 changes: 4 additions & 19 deletions imessage-exporter/src/app/runtime.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
The main app runtime.
*/


use std::{
cmp::min,
collections::{BTreeSet, HashMap, HashSet},
Expand Down Expand Up @@ -535,6 +534,7 @@ impl Config {
total_bytes: 100,
is_sticker: false,
hide_attachment: 0,
emoji_description: None,
copied_path: None,
}
}
Expand Down Expand Up @@ -921,23 +921,8 @@ mod who_tests {
#[cfg(test)]
mod directory_tests {
use crate::{Config, Options};
use imessage_database::tables::attachment::Attachment;
use std::path::PathBuf;

pub fn fake_attachment() -> Attachment {
Attachment {
rowid: 0,
filename: Some("a/b/c/d.jpg".to_string()),
uti: Some("public.png".to_string()),
mime_type: Some("image/png".to_string()),
transfer_name: Some("d.jpg".to_string()),
total_bytes: 100,
is_sticker: false,
hide_attachment: 0,
copied_path: None,
}
}

#[test]
fn can_get_valid_attachment_sub_dir() {
let options = Options::fake_options(crate::app::export_type::ExportType::Html);
Expand Down Expand Up @@ -983,7 +968,7 @@ mod directory_tests {
let app = Config::fake_app(options);

// Create attachment
let attachment = fake_attachment();
let attachment = Config::fake_attachment();

let result = app.message_attachment_path(&attachment);
let expected = String::from("a/b/c/d.jpg");
Expand All @@ -999,7 +984,7 @@ mod directory_tests {
let app = Config::fake_app(options);

// Create attachment
let mut attachment = fake_attachment();
let mut attachment = Config::fake_attachment();
let mut full_path = PathBuf::from("/Users/ReagentX/exports/attachments");
full_path.push(attachment.filename());
attachment.copied_path = Some(full_path);
Expand All @@ -1018,7 +1003,7 @@ mod directory_tests {
let app = Config::fake_app(options);

// Create attachment
let mut attachment = fake_attachment();
let mut attachment = Config::fake_attachment();
attachment.copied_path = Some(PathBuf::from(attachment.filename.as_ref().unwrap()));

let result = app.message_attachment_path(&attachment);
Expand Down
51 changes: 49 additions & 2 deletions imessage-exporter/src/exporters/html.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -595,14 +595,24 @@ impl<'a> Writer<'a> for HTML<'a> {

fn format_sticker(&self, sticker: &'a mut Attachment, message: &Message) -> String {
match self.format_attachment(sticker, message, &AttachmentMeta::default()) {
Ok(sticker_embed) => {
Ok(mut sticker_embed) => {
// Add sticker effect
let sticker_effect = sticker.get_sticker_effect(
&self.config.options.platform,
&self.config.options.db_path,
self.config.options.attachment_root.as_deref(),
);
if let Ok(Some(sticker_effect)) = sticker_effect {
return format!("{sticker_embed}\n<div class=\"sticker_effect\">Sent with {sticker_effect} effect</div>");
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_effect\">Sent with {sticker_effect} effect</div>"
))
}

// Add sticker prompt
if let Some(prompt) = &sticker.emoji_description {
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_effect\">Genmoji prompt: {prompt}</div>"
))
}
sticker_embed
}
Expand Down Expand Up @@ -2244,6 +2254,43 @@ mod tests {
std::fs::remove_file(orphaned_path).unwrap();
}

#[test]
fn can_format_html_attachment_sticker_genmoji() {
// Create exporter
let mut options = Options::fake_options(ExportType::Html);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();

let config = Config::fake_app(options);
let exporter = HTML::new(&config).unwrap();

let mut message = Config::fake_message();
// Set message to sticker variant
message.associated_message_type = Some(1000);

let mut attachment = Config::fake_attachment();
attachment.is_sticker = true;
attachment.emoji_description = Some("Example description".to_string());
let sticker_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/stickers/outline.heic");
attachment.filename = Some(sticker_path.to_string_lossy().to_string());
attachment.copied_path = Some(PathBuf::from(sticker_path.to_string_lossy().to_string()));

let actual = exporter.format_sticker(&mut attachment, &message);

assert_eq!(actual, "<img src=\"imessage-database/test_data/stickers/outline.heic\" loading=\"lazy\">\n<div class=\"sticker_effect\">Sent with Outline effect</div>\n<div class=\"sticker_effect\">Genmoji prompt: Example description</div>");

// Remove the file created by the constructor for this test
let orphaned_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("orphaned.html");
std::fs::remove_file(orphaned_path).unwrap();
}

#[test]
fn can_format_html_attachment_audio_transcript() {
// Create exporter
Expand Down
55 changes: 53 additions & 2 deletions imessage-exporter/src/exporters/txt.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -413,15 +413,24 @@ impl<'a> Writer<'a> for TXT<'a> {
);
match self.format_attachment(sticker, message, &AttachmentMeta::default()) {
Ok(path_to_sticker) => {
let mut out_s = format!("Sticker from {who}: {path_to_sticker}");

// Add sticker effect
let sticker_effect = sticker.get_sticker_effect(
&self.config.options.platform,
&self.config.options.db_path,
self.config.options.attachment_root.as_deref(),
);
if let Ok(Some(sticker_effect)) = sticker_effect {
return format!("{sticker_effect} Sticker from {who}: {path_to_sticker}");
out_s = format!("{sticker_effect} {out_s}");
}

// Add sticker prompt
if let Some(prompt) = &sticker.emoji_description {
out_s = format!("{out_s} (Genmoji prompt: {prompt})");
}
format!("Sticker from {who}: {path_to_sticker}")

out_s
}
Err(path) => format!("Sticker from {who}: {path}"),
}
Expand Down Expand Up @@ -1688,6 +1697,48 @@ mod tests {
std::fs::remove_file(orphaned_path).unwrap();
}

#[test]
fn can_format_txt_attachment_sticker_genmoji() {
// Create exporter
let mut options = Options::fake_options(ExportType::Txt);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();

let mut config = Config::fake_app(options);
config.participants.insert(0, ME.to_string());

let exporter = TXT::new(&config).unwrap();

let mut message = Config::fake_message();
// Set message to sticker variant
message.associated_message_type = Some(1000);

let mut attachment = Config::fake_attachment();
attachment.is_sticker = true;
attachment.emoji_description = Some("Example description".to_string());
let sticker_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/stickers/outline.heic");
attachment.filename = Some(sticker_path.to_string_lossy().to_string());
attachment.copied_path = Some(PathBuf::from(sticker_path.to_string_lossy().to_string()));

let actual = exporter.format_sticker(&mut attachment, &message);

assert_eq!(
actual,
"Outline Sticker from Me: imessage-database/test_data/stickers/outline.heic (Genmoji prompt: Example description)"
);

// Remove the file created by the constructor for this test
let orphaned_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("orphaned.txt");
std::fs::remove_file(orphaned_path).unwrap();
}

#[test]
fn can_format_txt_attachment_audio_transcript() {
// Create exporter
Expand Down

0 comments on commit d1c42af

Please sign in to comment.