From 4d04e51ddac9c0bda06ddd92619953f9417a50f7 Mon Sep 17 00:00:00 2001 From: KnugiHK <24708955+KnugiHK@users.noreply.github.com> Date: Sun, 2 Mar 2025 00:47:34 +0800 Subject: [PATCH] Refactor and add docstrings --- Whatsapp_Chat_Exporter/data_model.py | 168 +++++++++++++++++++-------- Whatsapp_Chat_Exporter/utility.py | 2 +- 2 files changed, 120 insertions(+), 50 deletions(-) diff --git a/Whatsapp_Chat_Exporter/data_model.py b/Whatsapp_Chat_Exporter/data_model.py index da14cc7..411de1c 100644 --- a/Whatsapp_Chat_Exporter/data_model.py +++ b/Whatsapp_Chat_Exporter/data_model.py @@ -1,37 +1,80 @@ -#!/usr/bin/python3 - import os from datetime import datetime, tzinfo, timedelta -from typing import Union, Optional +from typing import Union, Optional, Dict, Any -class Timing(): - def __init__(self, timezone_offset: Optional[int]): +class Timing: + """ + Handles timestamp formatting with timezone support. + """ + def __init__(self, timezone_offset: Optional[int]) -> None: + """ + Initialize Timing object. + + Args: + timezone_offset (Optional[int]): Hours offset from UTC + """ self.timezone_offset = timezone_offset - def format_timestamp(self, timestamp, format): + def format_timestamp(self, timestamp: Optional[Union[int, float]], format: str) -> Optional[str]: + """ + Format a timestamp with the specified format string. + + Args: + timestamp (Optional[Union[int, float]]): Unix timestamp to format + format (str): strftime format string + + Returns: + Optional[str]: Formatted timestamp string, or None if timestamp is None + """ if timestamp: timestamp = timestamp / 1000 if timestamp > 9999999999 else timestamp return datetime.fromtimestamp(timestamp, TimeZone(self.timezone_offset)).strftime(format) - else: - return None + return None class TimeZone(tzinfo): - def __init__(self, offset): + """ + Custom timezone class with fixed offset. + """ + def __init__(self, offset: int) -> None: + """ + Initialize TimeZone object. + + Args: + offset (int): Hours offset from UTC + """ self.offset = offset - def utcoffset(self, dt): - return timedelta(hours=self.offset) - def dst(self, dt): - return timedelta(0) + + def utcoffset(self, dt: Optional[datetime]) -> timedelta: + """Get UTC offset.""" + return timedelta(hours=self.offset) + + def dst(self, dt: Optional[datetime]) -> timedelta: + """Get DST offset (always 0).""" + return timedelta(0) -class ChatStore(): - def __init__(self, type, name=None, media=None): +class ChatStore: + """ + Stores chat information and messages. + """ + def __init__(self, type: str, name: Optional[str] = None, media: Optional[str] = None) -> None: + """ + Initialize ChatStore object. + + Args: + type (str): Device type (IOS or ANDROID) + name (Optional[str]): Chat name + media (Optional[str]): Path to media folder + + Raises: + TypeError: If name is not a string or None + """ if name is not None and not isinstance(name, str): raise TypeError("Name must be a string or None") self.name = name - self.messages = {} + self.messages: Dict[str, 'Message'] = {} self.type = type if media is not None: from Whatsapp_Chat_Exporter.utility import Device @@ -47,18 +90,20 @@ class ChatStore(): self.their_avatar_thumb = None self.status = None self.media_base = "" - - def add_message(self, id, message): + + def add_message(self, id: str, message: 'Message') -> None: + """Add a message to the chat store.""" if not isinstance(message, Message): raise TypeError("message must be a Message object") self.messages[id] = message - def delete_message(self, id): + def delete_message(self, id: str) -> None: + """Delete a message from the chat store.""" if id in self.messages: del self.messages[id] - def to_json(self): - serialized_msgs = {id: msg.to_json() for id, msg in self.messages.items()} + def to_json(self) -> Dict[str, Any]: + """Convert chat store to JSON-serializable dict.""" return { 'name': self.name, 'type': self.type, @@ -66,17 +111,22 @@ class ChatStore(): 'their_avatar': self.their_avatar, 'their_avatar_thumb': self.their_avatar_thumb, 'status': self.status, - 'messages': serialized_msgs + 'messages': {id: msg.to_json() for id, msg in self.messages.items()} } - def get_last_message(self): + def get_last_message(self) -> 'Message': + """Get the most recent message in the chat.""" return tuple(self.messages.values())[-1] - def get_messages(self): + def get_messages(self) -> Dict[str, 'Message']: + """Get all messages in the chat.""" return self.messages.values() -class Message(): +class Message: + """ + Represents a single message in a chat. + """ def __init__( self, *, @@ -87,17 +137,35 @@ class Message(): received_timestamp: int, read_timestamp: int, timezone_offset: int = 0, - message_type: int = None - ): + message_type: Optional[int] = None + ) -> None: + """ + Initialize Message object. + + Args: + from_me (Union[bool, int]): Whether message was sent by the user + timestamp (int): Message timestamp + time (Union[int, float, str]): Message time + key_id (int): Message unique identifier + received_timestamp (int): When message was received + read_timestamp (int): When message was read + timezone_offset (int, optional): Hours offset from UTC. Defaults to 0 + message_type (Optional[int], optional): Type of message. Defaults to None + + Raises: + TypeError: If time is not a string or number + """ self.from_me = bool(from_me) self.timestamp = timestamp / 1000 if timestamp > 9999999999 else timestamp timing = Timing(timezone_offset) - if isinstance(time, int) or isinstance(time, float): + + if isinstance(time, (int, float)): self.time = timing.format_timestamp(self.timestamp, "%H:%M") elif isinstance(time, str): self.time = time else: raise TypeError("Time must be a string or number") + self.media = False self.key_id = key_id self.meta = False @@ -107,29 +175,31 @@ class Message(): self.mime = None self.message_type = message_type, self.received_timestamp = timing.format_timestamp(received_timestamp, "%Y/%m/%d %H:%M") - self.read_timestamp = timing.format_timestamp(read_timestamp, "%Y/%m/%d %H:%M") - # Extra + self.read_timestamp = timing.format_timestamp(read_timestamp, "%Y/%m/%d %H:%M") + + # Extra attributes self.reply = None self.quoted_data = None self.caption = None - self.thumb = None # Android specific + self.thumb = None # Android specific self.sticker = False - - def to_json(self): + + def to_json(self) -> Dict[str, Any]: + """Convert message to JSON-serializable dict.""" return { - 'from_me' : self.from_me, - 'timestamp' : self.timestamp, - 'time' : self.time, - 'media' : self.media, - 'key_id' : self.key_id, - 'meta' : self.meta, - 'data' : self.data, - 'sender' : self.sender, - 'safe' : self.safe, - 'mime' : self.mime, - 'reply' : self.reply, - 'quoted_data' : self.quoted_data, - 'caption' : self.caption, - 'thumb' : self.thumb, - 'sticker' : self.sticker - } + 'from_me': self.from_me, + 'timestamp': self.timestamp, + 'time': self.time, + 'media': self.media, + 'key_id': self.key_id, + 'meta': self.meta, + 'data': self.data, + 'sender': self.sender, + 'safe': self.safe, + 'mime': self.mime, + 'reply': self.reply, + 'quoted_data': self.quoted_data, + 'caption': self.caption, + 'thumb': self.thumb, + 'sticker': self.sticker + } \ No newline at end of file diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py index 2db7445..b5c4e10 100644 --- a/Whatsapp_Chat_Exporter/utility.py +++ b/Whatsapp_Chat_Exporter/utility.py @@ -10,7 +10,7 @@ from markupsafe import Markup from datetime import datetime, timedelta from enum import IntEnum from Whatsapp_Chat_Exporter.data_model import ChatStore -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple try: from enum import StrEnum, IntEnum except ImportError: