diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py index 3f45365..609abca 100644 --- a/Whatsapp_Chat_Exporter/__main__.py +++ b/Whatsapp_Chat_Exporter/__main__.py @@ -16,6 +16,7 @@ from Whatsapp_Chat_Exporter.utility import APPLE_TIME, CLEAR_LINE, CURRENT_TZ_OF from Whatsapp_Chat_Exporter.utility import readable_to_bytes, safe_name, bytes_to_readable from Whatsapp_Chat_Exporter.utility import import_from_json, incremental_merge, check_update from Whatsapp_Chat_Exporter.utility import telegram_json_format, convert_time_unit, DbType +from Whatsapp_Chat_Exporter.utility import get_transcription_selection from argparse import ArgumentParser, SUPPRESS from datetime import datetime from getpass import getpass @@ -550,6 +551,7 @@ def process_messages(args, data: ChatCollection) -> None: # Process messages if args.android: message_handler = android_handler + data.set_system("transcription_selection", get_transcription_selection(db)) else: message_handler = ios_handler diff --git a/Whatsapp_Chat_Exporter/android_handler.py b/Whatsapp_Chat_Exporter/android_handler.py index 8d7dbad..eaf5704 100644 --- a/Whatsapp_Chat_Exporter/android_handler.py +++ b/Whatsapp_Chat_Exporter/android_handler.py @@ -82,7 +82,13 @@ def messages(db, data, media_folder, timezone_offset, filter_date, filter_chat, table_message = False except sqlite3.OperationalError: try: - content_cursor = _get_messages_cursor_new(c, filter_empty, filter_date, filter_chat) + content_cursor = _get_messages_cursor_new( + c, + filter_empty, + filter_date, + filter_chat, + data.get_system("transcription_selection"), + ) table_message = True except Exception as e: raise e @@ -217,7 +223,13 @@ def _get_messages_cursor_legacy(cursor, filter_empty, filter_date, filter_chat): return cursor -def _get_messages_cursor_new(cursor, filter_empty, filter_date, filter_chat): +def _get_messages_cursor_new( + cursor, + filter_empty, + filter_date, + filter_chat, + transcription_selection, + ): """Get cursor for new database schema.""" empty_filter = get_cond_for_empty(filter_empty, "key_remote_jid", "broadcast") date_filter = f'AND message.timestamp {filter_date}' if filter_date is not None else '' @@ -252,7 +264,7 @@ def _get_messages_cursor_new(cursor, filter_empty, filter_date, filter_chat): jid_global.type as jid_type, COALESCE(receipt_user.receipt_timestamp, message.received_timestamp) as received_timestamp, COALESCE(receipt_user.read_timestamp, receipt_user.played_timestamp) as read_timestamp, - message_media.raw_transcription_text as transcription_text + {transcription_selection} FROM message LEFT JOIN message_quoted ON message_quoted.message_row_id = message._id diff --git a/Whatsapp_Chat_Exporter/data_model.py b/Whatsapp_Chat_Exporter/data_model.py index 9ea0d7b..52f9bae 100644 --- a/Whatsapp_Chat_Exporter/data_model.py +++ b/Whatsapp_Chat_Exporter/data_model.py @@ -66,6 +66,7 @@ class ChatCollection(MutableMapping): def __init__(self) -> None: """Initialize an empty chat collection.""" self._chats: Dict[str, ChatStore] = {} + self._system: Dict[str, Any] = {} def __getitem__(self, key: str) -> 'ChatStore': """Get a chat by its ID. Required for dict-like access.""" @@ -148,6 +149,28 @@ class ChatCollection(MutableMapping): """ return {chat_id: chat.to_json() for chat_id, chat in self._chats.items()} + def get_system(self, key: str) -> Any: + """ + Get a system value by its key. + + Args: + key (str): The key of the system value to retrieve + + Returns: + Any: The system value if found, None otherwise + """ + return self._system.get(key) + + def set_system(self, key: str, value: Any) -> None: + """ + Set a system value by its key. + + Args: + key (str): The key of the system value to set + value (Any): The value to set + """ + self._system[key] = value + class ChatStore: """ diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py index 12c7da0..9e7047f 100644 --- a/Whatsapp_Chat_Exporter/utility.py +++ b/Whatsapp_Chat_Exporter/utility.py @@ -572,6 +572,24 @@ def get_status_location(output_folder: str, offline_static: str) -> str: return w3css +def get_transcription_selection(db: sqlite3.Connection) -> str: + """ + Returns the SQL selection statement for transcription text based on the database schema. + + Args: + db (sqlite3.Connection): The SQLite database connection. + Returns: + str: The SQL selection statement for transcription. + """ + cursor = db.cursor() + cursor.execute("PRAGMA table_info(message_media)") + columns = [row[1] for row in cursor.fetchall()] + + if "raw_transcription_text" in columns: + return "message_media.raw_transcription_text AS transcription_text" + else: + return "NULL AS transcription_text" + def setup_template(template: Optional[str], no_avatar: bool, experimental: bool = False) -> jinja2.Template: """ Sets up the Jinja2 template environment and loads the template.