tgdata 0.0.0__py3-none-any.whl

tgdata/__init__.py

@@ -0,0 +1,57 @@
+"""
+tgdata - Telegram Data Extraction Library
+
+A production-grade Python library for extracting and processing Telegram group and channel messages.
+"""
+
+# Main class
+from .tgdata import TgData
+
+# Models
+from .models import (
+    MessageData,
+    GroupInfo,
+    ConnectionConfig,
+    RateLimitInfo
+)
+
+# Utilities
+from .utils import (
+    format_message_for_display,
+    export_to_json,
+    export_to_csv,
+    filter_messages_by_sender,
+    filter_messages_by_content,
+    get_message_statistics,
+    save_profile_photos,
+    create_metrics_report
+)
+
+# Progress tracking
+from .progress import ProgressTracker
+
+__version__ = "0.0.0"
+
+__all__ = [
+    # Main class
+    "TgData",
+    
+    # Models
+    "MessageData",
+    "GroupInfo",
+    "ConnectionConfig",
+    "RateLimitInfo",
+    
+    # Utilities
+    "format_message_for_display",
+    "export_to_json",
+    "export_to_csv",
+    "filter_messages_by_sender",
+    "filter_messages_by_content",
+    "get_message_statistics",
+    "save_profile_photos",
+    "create_metrics_report",
+    
+    # Progress
+    "ProgressTracker"
+]

tgdata/connection_engine.py

@@ -0,0 +1,316 @@
+"""
+Connection management engine for Telegram.
+Handles connection pooling, rate limiting, retries, and health checks.
+"""
+
+import asyncio
+import configparser
+import logging
+import time
+from typing import Optional, List, Dict, Any, Union
+from telethon import TelegramClient
+from telethon.errors import FloodWaitError, AuthKeyUnregisteredError
+from telethon.sessions import StringSession
+
+from .models import ConnectionConfig, RateLimitInfo
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionPool:
+    """Manage multiple connections with rate limiting"""
+    
+    def __init__(self, max_connections: int = 3):
+        self.max_connections = max_connections
+        self.connections: List[TelegramClient] = []
+        self.current_index = 0
+        self.rate_limits: Dict[int, RateLimitInfo] = {}
+        
+    async def get_connection(self) -> TelegramClient:
+        """Get next available connection using round-robin"""
+        if not self.connections:
+            raise ValueError("No connections in pool")
+            
+        # Try to find a connection not rate-limited
+        attempts = 0
+        while attempts < len(self.connections):
+            conn = self.connections[self.current_index]
+            self.current_index = (self.current_index + 1) % len(self.connections)
+            
+            rate_info = self.rate_limits.get(id(conn), RateLimitInfo())
+            if time.time() < rate_info.flood_wait_until:
+                attempts += 1
+                continue
+                
+            return conn
+            
+        # All connections rate-limited, return the one with shortest wait
+        return min(self.connections, 
+                  key=lambda c: self.rate_limits.get(id(c), RateLimitInfo()).flood_wait_until)
+        
+    def mark_rate_limited(self, conn: TelegramClient, wait_seconds: float):
+        """Mark a connection as rate-limited"""
+        rate_info = self.rate_limits.get(id(conn), RateLimitInfo())
+        rate_info.flood_wait_until = time.time() + wait_seconds
+        self.rate_limits[id(conn)] = rate_info
+        logger.warning(f"Connection {id(conn)} rate-limited for {wait_seconds}s")
+    
+    def add_connection(self, conn: TelegramClient):
+        """Add a connection to the pool"""
+        self.connections.append(conn)
+        
+    async def close_all(self):
+        """Close all connections in the pool"""
+        for conn in self.connections:
+            await conn.disconnect()
+        self.connections.clear()
+
+
+class ConnectionEngine:
+    """
+    Manages Telegram connections with advanced features.
+    """
+    
+    def __init__(self, 
+                 config_path: str = "config.ini",
+                 pool_size: int = 1,
+                 max_retries: int = 3,
+                 retry_delay: float = 1.0,
+                 exponential_backoff: bool = True):
+        """
+        Initialize connection engine.
+        
+        Args:
+            config_path: Path to configuration file
+            pool_size: Number of connections in pool (1 = no pooling)
+            max_retries: Maximum retry attempts
+            retry_delay: Initial retry delay in seconds
+            exponential_backoff: Whether to use exponential backoff
+        """
+        self.config_path = config_path
+        self.pool_size = pool_size
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self.exponential_backoff = exponential_backoff
+        
+        self._config: Optional[ConnectionConfig] = None
+        self._primary_client: Optional[TelegramClient] = None
+        self._pool: Optional[ConnectionPool] = None
+        self._health_check_interval = 300  # 5 minutes
+        self._last_health_check = 0
+        
+    def _load_config(self) -> ConnectionConfig:
+        """Load configuration from file"""
+        if self._config:
+            return self._config
+            
+        config = configparser.ConfigParser()
+        config.read(self.config_path)
+        
+        # Handle case-insensitive section names
+        telegram_section = None
+        for section in config.sections():
+            if section.lower() == 'telegram':
+                telegram_section = section
+                break
+        
+        if telegram_section is None:
+            raise ValueError("No [telegram] or [Telegram] section found in config file")
+        
+        self._config = ConnectionConfig(
+            api_id=config[telegram_section]['api_id'],
+            api_hash=config[telegram_section]['api_hash'],
+            session_file=config[telegram_section].get('session_file', 'telegram_session'),
+            phone=config[telegram_section].get('phone'),
+            username=config[telegram_section].get('username'),
+            max_retries=self.max_retries,
+            retry_delay=self.retry_delay,
+            exponential_backoff=self.exponential_backoff
+        )
+        
+        return self._config
+        
+    async def get_client(self) -> TelegramClient:
+        """
+        Get a connected Telegram client.
+        Uses pooling if enabled, otherwise returns primary client.
+        """
+        # Check if health check needed
+        current_time = time.time()
+        if current_time - self._last_health_check > self._health_check_interval:
+            await self.health_check()
+            self._last_health_check = current_time
+            
+        # Use pool if enabled
+        if self.pool_size > 1 and self._pool:
+            return await self._pool.get_connection()
+            
+        # Otherwise ensure primary client is connected
+        if not self._primary_client:
+            await self._init_primary_client()
+            
+        if not self._primary_client.is_connected():
+            await self._connect_with_retry(self._primary_client)
+            
+        return self._primary_client
+        
+    async def _init_primary_client(self):
+        """Initialize the primary client"""
+        config = self._load_config()
+        
+        # Use username as session name for compatibility with original code
+        session_name = config.username if config.username else config.session_file
+        self._primary_client = TelegramClient(
+            session_name,
+            config.api_id,
+            config.api_hash
+        )
+        
+        await self._connect_with_retry(self._primary_client)
+        
+        # Initialize pool if needed
+        if self.pool_size > 1:
+            await self._init_pool()
+            
+    async def _init_pool(self):
+        """Initialize connection pool"""
+        config = self._load_config()
+        self._pool = ConnectionPool(self.pool_size)
+        
+        # Add primary client to pool
+        if self._primary_client:
+            self._pool.add_connection(self._primary_client)
+            
+        # Create additional connections
+        for i in range(1, self.pool_size):
+            session_file = f"{config.session_file}_{i}"
+            client = TelegramClient(
+                session_file,
+                config.api_id,
+                config.api_hash
+            )
+            
+            await self._connect_with_retry(client)
+            self._pool.add_connection(client)
+            
+        logger.info(f"Initialized connection pool with {self.pool_size} connections")
+        
+    async def _connect_with_retry(self, client: TelegramClient):
+        """Connect a client with retry logic"""
+        config = self._load_config()
+        
+        try:
+            # Use start() directly - it handles everything including retries
+            await client.start(phone=config.phone)
+            logger.info("Successfully connected and authenticated!")
+            
+        except FloodWaitError as e:
+            wait_time = e.seconds
+            logger.warning(f"Rate limited, waiting {wait_time} seconds...")
+            
+            if self._pool:
+                self._pool.mark_rate_limited(client, wait_time)
+                
+            await asyncio.sleep(wait_time)
+            # Retry after rate limit
+            await client.start(phone=config.phone)
+            
+        except Exception as e:
+            logger.error(f"Connection failed: {e}")
+            raise ConnectionError(f"Failed to connect: {e}")
+                    
+    async def health_check(self) -> Dict[str, Any]:
+        """
+        Perform health check on all connections.
+        
+        Returns:
+            Dictionary with health status
+        """
+        logger.info("Performing connection health check...")
+        
+        status = {
+            'timestamp': time.time(),
+            'primary_connection': False,
+            'pool_connections': [],
+            'errors': []
+        }
+        
+        try:
+            # Check primary connection
+            if self._primary_client:
+                try:
+                    if self._primary_client.is_connected():
+                        await self._primary_client.get_me()
+                        status['primary_connection'] = True
+                        logger.info("Primary connection healthy")
+                except Exception as e:
+                    status['errors'].append(f"Primary connection error: {e}")
+                    logger.warning(f"Primary connection unhealthy: {e}")
+                    
+            # Check pool connections
+            if self._pool:
+                for i, conn in enumerate(self._pool.connections):
+                    try:
+                        if conn.is_connected():
+                            await conn.get_me()
+                            status['pool_connections'].append({
+                                'index': i,
+                                'healthy': True,
+                                'rate_limited': time.time() < self._pool.rate_limits.get(
+                                    id(conn), RateLimitInfo()
+                                ).flood_wait_until
+                            })
+                        else:
+                            status['pool_connections'].append({
+                                'index': i,
+                                'healthy': False
+                            })
+                    except Exception as e:
+                        status['errors'].append(f"Pool connection {i} error: {e}")
+                        status['pool_connections'].append({
+                            'index': i,
+                            'healthy': False,
+                            'error': str(e)
+                        })
+                        
+        except Exception as e:
+            status['errors'].append(f"Health check error: {e}")
+            logger.error(f"Health check failed: {e}")
+            
+        return status
+        
+    async def validate_connection(self) -> bool:
+        """
+        Validate that we have a working connection.
+        
+        Returns:
+            True if connection is valid
+        """
+        try:
+            client = await self.get_client()
+            await client.get_me()
+            return True
+        except Exception as e:
+            logger.error(f"Connection validation failed: {e}")
+            return False
+            
+    async def handle_rate_limit(self, error: FloodWaitError, client: Optional[TelegramClient] = None):
+        """Handle rate limit errors"""
+        wait_time = error.seconds
+        logger.warning(f"Rate limit hit! Waiting {wait_time} seconds...")
+        
+        if self._pool and client:
+            self._pool.mark_rate_limited(client, wait_time)
+            
+        await asyncio.sleep(wait_time)
+        
+    async def close(self):
+        """Close all connections"""
+        if self._pool:
+            await self._pool.close_all()
+        elif self._primary_client:
+            await self._primary_client.disconnect()
+            
+        self._primary_client = None
+        self._pool = None
+        logger.info("All connections closed")

tgdata/message_engine.py

@@ -0,0 +1,299 @@
+"""
+Message processing engine for Telegram.
+Handles message fetching, filtering, deduplication, and formatting.
+"""
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Optional, List, Dict, Any, Callable
+import pandas as pd
+from telethon.errors import FloodWaitError
+
+from .connection_engine import ConnectionEngine
+from .models import MessageData
+from .progress import ProgressTracker
+
+logger = logging.getLogger(__name__)
+
+
+class MessageEngine:
+    """
+    Handles all message-related operations.
+    """
+    
+    def __init__(self,
+                 connection_engine: ConnectionEngine):
+        """
+        Initialize message engine.
+        
+        Args:
+            connection_engine: Connection engine instance
+        """
+        self.connection_engine = connection_engine
+            
+        # Message cache for performance
+        self._message_cache: Dict[str, pd.DataFrame] = {}
+        
+    async def fetch_messages(self,
+                           group_id: int,
+                           limit: Optional[int] = None,
+                           start_date: Optional[datetime] = None,
+                           end_date: Optional[datetime] = None,
+                           include_profile_photos: bool = False,
+                           use_cache: bool = True,
+                           progress_callback: Optional[Callable] = None,
+                           min_id: Optional[int] = None) -> pd.DataFrame:
+        """
+        Fetch messages from a group with various filters.
+        
+        Args:
+            group_id: Telegram group/channel ID
+            limit: Maximum number of messages to fetch
+            start_date: Get messages after this date
+            end_date: Get messages before this date
+            include_profile_photos: Whether to download profile photos
+            use_cache: Whether to use cached results
+            progress_callback: Optional callback for progress updates
+            min_id: Minimum message ID (for resuming)
+            
+        Returns:
+            DataFrame with messages
+        """
+        # Check cache
+        cache_key = f"{group_id}_{start_date}_{end_date}_{limit}"
+        if use_cache and cache_key in self._message_cache:
+            logger.info("Returning cached messages")
+            return self._message_cache[cache_key]
+            
+        # Set up progress tracking
+        progress_tracker = None
+        if progress_callback:
+            progress_tracker = ProgressTracker(
+                total_expected=limit,
+                callback=progress_callback
+            )
+            progress_tracker.start()
+            
+        messages_data = []
+        processed_count = 0
+        
+        try:
+            client = await self.connection_engine.get_client()
+            
+            async with client:
+                # Get the entity
+                channel = await client.get_entity(group_id)
+                logger.info(f"Fetching messages from: {channel.title}")
+                
+                # Determine iteration parameters
+                offset_date = end_date if end_date else datetime.now()
+                reverse = bool(start_date)  # Reverse if we have a start date
+                
+                # Build kwargs for iter_messages
+                iter_kwargs = {
+                    'entity': channel,
+                    'limit': limit,
+                    'offset_date': offset_date,
+                    'reverse': reverse
+                }
+                
+                # Only add min_id if it's not None
+                if min_id is not None:
+                    iter_kwargs['min_id'] = min_id
+                
+                # Iterate through messages
+                async for msg in client.iter_messages(**iter_kwargs):
+                    # Apply date filters
+                    if start_date and msg.date < start_date:
+                        continue
+                    if end_date and msg.date > end_date:
+                        break
+                        
+                    # Check deduplication
+                    if self.enable_deduplication and self.tracker:
+                        if await self.tracker.is_processed(msg.id, group_id):
+                            logger.debug(f"Skipping duplicate message {msg.id}")
+                            continue
+                            
+                    # Process message
+                    message_data = await self._process_message(
+                        msg, 
+                        client,
+                        include_profile_photos
+                    )
+                    
+                    if message_data:
+                        messages_data.append(message_data.to_dict())
+                        
+                            
+                        processed_count += 1
+                        
+                        # Update progress
+                        if progress_tracker:
+                            progress_tracker.update()
+                            
+                        # Log progress
+                        if processed_count % 100 == 0:
+                            logger.info(f"Processed {processed_count} messages...")
+                            
+        except FloodWaitError as e:
+            await self.connection_engine.handle_rate_limit(e, client)
+            # Retry after rate limit
+            return await self.fetch_messages(
+                group_id=group_id,
+                limit=limit,
+                start_date=start_date,
+                end_date=end_date,
+                include_profile_photos=include_profile_photos,
+                use_cache=use_cache,
+                progress_callback=progress_callback,
+                min_id=min_id
+            )
+            
+        except Exception as e:
+            logger.error(f"Error fetching messages: {e}")
+            raise
+            
+        # Create DataFrame
+        df = pd.DataFrame(messages_data)
+        
+        # Cache results
+        if use_cache and not df.empty:
+            self._message_cache[cache_key] = df
+            
+        logger.info(f"Retrieved {len(df)} messages")
+        
+        # Log tracker stats
+        if self.tracker:
+            stats = await self.tracker.get_stats(group_id)
+            logger.info(f"Tracker stats: {stats}")
+            
+        return df
+        
+    async def _process_message(self,
+                             msg,
+                             client,
+                             include_profile_photos: bool) -> Optional[MessageData]:
+        """Process a single message"""
+        try:
+            sender = await msg.get_sender()
+            if not sender:
+                return None
+                
+            # Extract sender info
+            sender_name = f"{getattr(sender, 'first_name', '') or ''} {getattr(sender, 'last_name', '') or ''}".strip()
+            username = f"@{sender.username}" if hasattr(sender, 'username') and sender.username else "No username"
+            
+            # Create message data
+            message_data = MessageData(
+                message_id=msg.id,
+                sender_id=sender.id,
+                sender_name=sender_name,
+                username=username,
+                message=msg.message,
+                date=msg.date,
+                reply_to_id=msg.reply_to_msg_id,
+                forwarded_from=msg.fwd_from.from_id if msg.fwd_from else None
+            )
+            
+            # Download profile photo if requested
+            if include_profile_photos and hasattr(sender, 'photo') and sender.photo:
+                try:
+                    photo_bytes = await client.download_profile_photo(sender, file=bytes)
+                    message_data.photo_data = photo_bytes
+                except Exception as e:
+                    logger.warning(f"Failed to download photo for {sender.id}: {e}")
+                    
+            return message_data
+            
+        except Exception as e:
+            logger.error(f"Error processing message {msg.id}: {e}")
+            return None
+            
+        
+    def clear_cache(self, group_id: Optional[int] = None):
+        """
+        Clear message cache.
+        
+        Args:
+            group_id: Optional group ID to clear cache for specific group
+        """
+        if group_id:
+            # Clear cache for specific group
+            keys_to_remove = [k for k in self._message_cache.keys() if k.startswith(f"{group_id}_")]
+            for key in keys_to_remove:
+                del self._message_cache[key]
+            logger.info(f"Cleared cache for group {group_id}")
+        else:
+            # Clear all cache
+            self._message_cache.clear()
+            logger.info("Cleared all message cache")
+            
+    async def get_message_count(self, group_id: int) -> int:
+        """
+        Get total message count for a group.
+        
+        Args:
+            group_id: Telegram group/channel ID
+            
+        Returns:
+            Total message count
+        """
+        try:
+            client = await self.connection_engine.get_client()
+            
+            async with client:
+                # Get the messages with limit=1 to access count
+                messages = await client.get_messages(group_id, limit=1)
+                
+                # TotalList objects have a 'total' attribute
+                if hasattr(messages, 'total'):
+                    return messages.total
+                else:
+                    # If it's just a regular list, we can't get total efficiently
+                    logger.warning("Could not get message count efficiently")
+                    return 0
+                    
+        except Exception as e:
+            logger.error(f"Error getting message count: {e}")
+            raise
+            
+    async def search_messages(self,
+                            group_id: int,
+                            query: str,
+                            limit: Optional[int] = None) -> pd.DataFrame:
+        """
+        Search for messages containing specific text.
+        
+        Args:
+            group_id: Telegram group/channel ID
+            query: Search query
+            limit: Maximum number of results
+            
+        Returns:
+            DataFrame with matching messages
+        """
+        try:
+            client = await self.connection_engine.get_client()
+            messages_data = []
+            
+            async with client:
+                channel = await client.get_entity(group_id)
+                
+                async for msg in client.iter_messages(
+                    channel,
+                    search=query,
+                    limit=limit
+                ):
+                    message_data = await self._process_message(msg, client, False)
+                    if message_data:
+                        messages_data.append(message_data.to_dict())
+                        
+            df = pd.DataFrame(messages_data)
+            logger.info(f"Found {len(df)} messages matching '{query}'")
+            return df
+            
+        except Exception as e:
+            logger.error(f"Error searching messages: {e}")
+            raise