Refactor and add docstrings

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
+        }

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: