add types support

api_server.py

@@ -12,7 +12,7 @@
 import logging
 import os
 import mimetypes
-from typing import List
+from typing import List, Union
 
 import json
 from datetime import datetime
@@ -197,10 +197,10 @@ async def prepare_file_response(file_path: str, delete_after: bool = False):
     else:
         return FileResponse(path=file_path, media_type=media_type, headers=headers)
 
-async def download_media_file(channel: str, post_id: int, file_unique_id: str) -> str:
+async def download_media_file(channel: Union[str, int], post_id: int, file_unique_id: str) -> tuple[Union[str, None], bool]:
     """
     Download media file from Telegram and save to cache
-    Returns path to downloaded file
+    Returns tuple of (file path, delete_after)
     """
     base_cache_dir = os.path.abspath("./data/cache")
     
@@ -210,7 +210,7 @@ async def download_media_file(channel: str, post_id: int, file_unique_id: str) -
     os.makedirs(post_dir, exist_ok=True)
     
     # Convert numeric channel ID to int if needed
-    channel_id = channel
+    channel_id: Union[str, int] = channel
     if isinstance(channel, str) and channel.startswith('-100'):
         channel_id = int(channel)
     
@@ -467,9 +467,10 @@ def fix_corrupted_json(file_path: str) -> list:
         
         # Validate entries
         valid_entries = []
-        for entry in fixed_data:
-            if isinstance(entry, dict) and all(key in entry for key in ['channel', 'post_id', 'file_unique_id']):
-                valid_entries.append(entry)
+        if isinstance(fixed_data, list):
+            for entry in fixed_data:
+                if isinstance(entry, dict) and all(key in entry for key in ['channel', 'post_id', 'file_unique_id']):
+                    valid_entries.append(entry)
                 
         logger.info(f"Fixed JSON file: {file_path}, found {len(valid_entries)} valid entries")
         return valid_entries
@@ -584,13 +585,19 @@ def calculate_cache_stats():
         "channels": channels_stats
     }
 
+def is_local_request(request: Union[Request, None]) -> bool:
+    local_hosts = ["127.0.0.1", "localhost"]
+    if request and request.client and request.client.host:
+        if request.client.host in local_hosts:
+            return True
+    return False
 
 @app.get("/html/{channel}/{post_id}", response_class=HTMLResponse)
 @app.get("/post/html/{channel}/{post_id}", response_class=HTMLResponse)
 @app.get("/html/{channel}/{post_id}/{token}", response_class=HTMLResponse)  
 @app.get("/post/html/{channel}/{post_id}/{token}", response_class=HTMLResponse)
-async def get_post_html(channel: str, post_id: int, token: str | None = None, debug: bool = False, request: Request = None):
-    if Config["token"] and request and request.client.host not in ["127.0.0.1", "localhost"]:
+async def get_post_html(channel: str, post_id: int, request: Request, token: str | None = None, debug: bool = False):
+    if Config["token"] and is_local_request(request):
         if token != Config["token"]:
             logger.error(f"Invalid token for HTML post: {token}, expected: {Config['token']}")
             raise HTTPException(status_code=403, detail="Invalid token")
@@ -613,8 +620,8 @@ async def get_post_html(channel: str, post_id: int, token: str | None = None, de
 @app.get("/post/json/{channel}/{post_id}")
 @app.get("/json/{channel}/{post_id}/{token}")
 @app.get("/post/json/{channel}/{post_id}/{token}")
-async def get_post(channel: str, post_id: int, token: str | None = None, debug: bool = False, request: Request = None):
-    if Config["token"] and request and request.client.host not in ["127.0.0.1", "localhost"]:
+async def get_post(channel: str, post_id: int, request: Request, token: str | None = None, debug: bool = False):
+    if Config["token"] and is_local_request(request):
         if token != Config["token"]:
             logger.error(f"Invalid token for JSON post: {token}, expected: {Config['token']}")
             raise HTTPException(status_code=403, detail="Invalid token")
@@ -635,8 +642,8 @@ async def get_post(channel: str, post_id: int, token: str | None = None, debug:
 
 @app.get("/raw_json/{channel}/{post_id}")
 @app.get("/raw_json/{channel}/{post_id}/{token}")
-async def get_raw_post_json(channel: str, post_id: int, token: str | None = None, request: Request = None):
-    if Config["token"] and request and request.client.host not in ["127.0.0.1", "localhost"]:
+async def get_raw_post_json(channel: str, post_id: int, request: Request, token: str | None = None):
+    if Config["token"] and is_local_request(request):
         if token != Config["token"]:
             logger.error(f"Invalid token for raw JSON post: {token}, expected: {Config['token']}")
             raise HTTPException(status_code=403, detail="Invalid token")
@@ -645,7 +652,7 @@ async def get_raw_post_json(channel: str, post_id: int, token: str | None = None
             
     try:
         # Convert numeric channel ID to int if needed
-        channel_id = channel
+        channel_id: Union[str, int] = channel
         if isinstance(channel, str) and channel.startswith('-100'):
             channel_id = int(channel)
             
@@ -663,8 +670,8 @@ async def get_raw_post_json(channel: str, post_id: int, token: str | None = None
 
 @app.get("/health")
 @app.get("/health/{token}")
-async def health_check(token: str | None = None, request: Request = None):
-    if Config["token"] and request and request.client.host not in ["127.0.0.1", "localhost"]:
+async def health_check(request: Request, token: str | None = None):
+    if Config["token"] and is_local_request(request):
         if token != Config["token"]:
             logger.error(f"Invalid token for health check: {token}, expected: {Config['token']}")
             raise HTTPException(status_code=403, detail="Invalid token")
@@ -712,13 +719,16 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, digest: str
             logger.info(f"Valid digest for media {url}: {digest}")   
             
         # Convert numeric channel ID to int if needed
-        channel_id = channel
+        channel_id: Union[str, int] = channel
         if isinstance(channel, str) and channel.startswith('-100'):
             channel_id = int(channel)
             
         try: # Wrap the download and prepare call
             file_path, delete_after = await download_media_file(channel_id, post_id, file_unique_id)
-            return await prepare_file_response(file_path, delete_after=delete_after)
+            if not file_path:
+                raise HTTPException(status_code=404, detail="File not found")
+            if file_path:
+                return await prepare_file_response(file_path, delete_after=delete_after)
         except ZeroSizeFileError as e: # Catch zero-size file errors
             logger.warning(f"zero_size_file_encountered: {str(e)}. Instructing client to retry.")
             return Response(
@@ -744,15 +754,15 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, digest: str
 @app.get("/rss/{channel}", response_class=Response)
 @app.get("/rss/{channel}/{token}", response_class=Response)
 async def get_rss_feed(channel: str, 
+                        request: Request,
                         token: str | None = None, 
                         limit: int = 50, 
                         output_type: str = 'rss', 
                         exclude_flags: str | None = None,
                         exclude_text: str | None = None,
                         merge_seconds: int = 5,
-                        request: Request = None
                         ):
-    if Config["token"] and request and request.client.host not in ["127.0.0.1", "localhost"]:
+    if Config["token"] and is_local_request(request):
         if token != Config["token"]:
             logger.error(f"invalid_token_error: token {token}, expected {Config['token']}")
             raise HTTPException(status_code=403, detail="Invalid token")
@@ -798,9 +808,9 @@ async def get_rss_feed(channel: str,
 
 @app.get("/flags", response_model=List[str])
 @app.get("/flags/{token}", response_model=List[str])
-async def get_available_flags(token: str | None = None, request: Request = None):
+async def get_available_flags(request: Request, token: str | None = None):
     """Returns a list of all possible flags that can be assigned to posts."""
-    if Config["token"] and request and request.client.host not in ["127.0.0.1", "localhost"]:
+    if Config["token"] and is_local_request(request):
         if token != Config["token"]:
             logger.error(f"Invalid token for flags endpoint: {token}, expected: {Config['token']}")
             raise HTTPException(status_code=403, detail="Invalid token")

post_parser.py

@@ -90,7 +90,7 @@ class PostParser:
         else:
             return channel
 
-    async def get_post(self, channel: str, post_id: int, output_type: str = 'json', debug: bool = False) -> Union[str, Dict[Any, Any]]:
+    async def get_post(self, channel: str, post_id: int, output_type: str = 'json', debug: bool = False) -> Union[str, Dict[Any, Any], None]:
         print(f"Getting post {channel}, {post_id}")
         try:
             channel = self.channel_name_prepare(channel)
@@ -168,7 +168,7 @@ class PostParser:
             else: return title_segment
         else: return first_line
         
-    def _service_message_title(self, message: Message) -> str:
+    def _service_message_title(self, message: Message) -> Union[str, None]:
         if service := getattr(message, "service", None):
             if 'PINNED_MESSAGE'           in str(service): return "📌 Pinned message"
             elif 'NEW_CHAT_PHOTO'         in str(service): return "🖼 New chat photo"
@@ -179,17 +179,18 @@ class PostParser:
             elif 'GROUP_CHAT_CREATED'     in str(service): return "✨ Group chat created"
             elif 'CHANNEL_CHAT_CREATED'   in str(service): return "✨ Chat created"
             elif 'DELETE_CHAT_PHOTO'      in str(service): return "🗑️ Chat photo deleted"
+        return None
 
-    def _media_message_title(self, message: Message) -> str:
+    def _media_message_title(self, message: Message) -> Union[str, None]:
         if message.media:
             if message.media == MessageMediaType.POLL:
-                if hasattr(message, 'poll') and hasattr(message.poll, 'question'):
+                if message.poll is not None and hasattr(message.poll, 'question'):
                     poll_question = message.poll.question.strip()
                     if poll_question:
                         return f"📊 Poll: {poll_question}"
                 return "📊 Poll"
             if message.media == MessageMediaType.DOCUMENT:
-                if hasattr(message.document, 'mime_type') and 'pdf' in message.document.mime_type.lower():
+                if message.document is not None and hasattr(message.document, 'mime_type') and message.document.mime_type == 'application/pdf':
                     return "📄 PDF Document"
                 return "📎 Document"
             if message.media == MessageMediaType.PHOTO:       return "📷 Photo"
@@ -205,9 +206,10 @@ class PostParser:
             if message.web_page.title:
                 return f"🔗 {message.web_page.title}"
             return "🔗 Web link"
+        return None
         
 
-    def _generate_base_title(self, message: Message) -> str:
+    def _generate_base_title(self, message: Message) -> Union[str, None]:
         """Generates the base title without the FWD prefix."""
         # --- Text Processing --- (Phase 1: Process text if available)
         text = message.text or message.caption or ''
@@ -260,7 +262,8 @@ class PostParser:
                 if url_match: return f"🔗 {message.web_page.title}"
             if text_has_urls: # If original text had any URL (and wasn't YouTube/Webpage with title)
                 return  "🔗 Web link"
-
+        return None
+    
     def _generate_title(self, message: Message) -> str: #Tests: tests/postparser_gen_title.py
         """Generate a title for a message, based on its content."""
         title = None
@@ -304,7 +307,7 @@ class PostParser:
         
         return None
     
-    def _extract_reactions(self, message: Message) -> Union[str, None]:
+    def _extract_reactions(self, message: Message) -> Union[Dict[str, int], None]:
         if reactions := getattr(message, 'reactions', None):
             return {r.emoji: r.count for r in reactions.reactions}
         return None
@@ -360,8 +363,8 @@ class PostParser:
             flags.append("donat")
 
         # Check if the post's reactions contain more clown emojis (🤡) or poo emojis (💩).
-        if getattr(message, "reactions", None):
-            for reaction in message.reactions.reactions:
+        if reactions := getattr(message, "reactions", None):
+            for reaction in getattr(reactions, "reactions", []):
                 if reaction.emoji == "🤡" and reaction.count >= 30:
                     flags.append("clownpoo")
                     break
@@ -450,8 +453,8 @@ class PostParser:
         # Add raw JSON debug output if debug is enabled
         if debug:
             html_content.append(f'<pre class="debug-json" style="background: #f5f5f5; padding: 10px; margin-top: 20px; overflow-x: auto; font-size: 10px; white-space: pre-wrap;">{str(message)}</pre>')
-        html_content = '\n'.join(html_content)
-        return html_content
+        html_data = '\n'.join(html_content)
+        return html_data
 
     def _format_flags(self, flags_list: list) -> str:
         if not Config['show_post_flags']:
@@ -506,11 +509,11 @@ class PostParser:
         allowed_tags = ['p', 'a', 'b', 'i', 'strong', 'em', 'ul', 'ol', 'li', 'br', 'div', 'span', 'img', 'video', 'audio', 'source']
         allowed_attributes = {
             'a': ['href', 'title', 'target'],
-            'img': {'src': True, 'alt': True, 'style': True},
-            'video': {'controls': True, 'src': True, 'style': True},
-            'audio': {'controls': True, 'style': True},
+            'img': ['src', 'alt', 'style'],
+            'video': ['controls', 'src', 'style'],
+            'audio': ['controls', 'style'],
             'source': ['src', 'type'],
-            'div': {'class': True, 'style': True},
+            'div': ['class', 'style'],
             'span': ['class']
         }
 
@@ -551,9 +554,11 @@ class PostParser:
         if poll := getattr(message, "poll", None):
             poll_html = self._format_poll(poll)
 
+        forward_html = self._format_forward_info(message)
+
         if text_html or poll_html: 
             content_body.append(f'<div class="message-text">')
-            if message.forward_origin: content_body.append(self._format_forward_info(message)) # Forward info
+            if forward_html: content_body.append(forward_html) # Forward info
             content_body.append(f'{text_html}')
             if poll_html: content_body.append(poll_html) # Poll
             if message.forward_origin: content_body.append(f"<br>---- Forward post end ----") # Forward info end
@@ -582,7 +587,7 @@ class PostParser:
                 content_media.append(f'<div class="message-media">')
                 
                 # Check if document is a PDF file
-                if message.media == MessageMediaType.DOCUMENT and hasattr(message.document, 'mime_type') and message.document.mime_type == 'application/pdf':
+                if message.media == MessageMediaType.DOCUMENT and message.document is not None and hasattr(message.document, 'mime_type') and message.document.mime_type == 'application/pdf':
                     if channel_username.startswith('-100'): tg_link = f"https://t.me/c/{channel_username[4:]}/{message.id}"
                     else:  tg_link = f"https://t.me/{channel_username}/{message.id}"
                     content_media.append(f'<div class="document-pdf" style="padding: 10px;"><a href="{tg_link}" target="_blank">[PDF-файл]</a></div>')
@@ -829,10 +834,10 @@ class PostParser:
     def _save_media_file_ids(self, message: Message) -> None:
         try:
             file_data = {
-                'channel': None,
-                'post_id': None,
-                'file_unique_id': None,
-                'added': None
+                'channel': '',
+                'post_id': 0,
+                'file_unique_id': '',
+                'added': .0
             }
 
             channel_username = self.get_channel_username(message)
@@ -842,7 +847,7 @@ class PostParser:
 
             if message.media:
                 # Skip large videos - they shouldn't be cached permanently
-                if message.video and message.video.file_size > 100 * 1024 * 1024: return
+                if message.video and message.video.file_size and message.video.file_size > 100 * 1024 * 1024: return
                 
                 if   message.photo:         file_data['file_unique_id'] = message.photo.file_unique_id
                 elif message.video:         file_data['file_unique_id'] = message.video.file_unique_id

requirements.txt

@@ -10,4 +10,5 @@ feedgen==1.0.0
 python-dateutil==2.9.0.post0
 python-magic==0.4.27
 bleach[css]==6.1.0
+types-bleach
 json-repair

rss_generator.py

@@ -7,6 +7,7 @@
 # pylint: disable=broad-exception-caught, missing-function-docstring, missing-class-docstring
 # pylint: disable=f-string-without-interpolation
 # pylance: disable=reportMissingImports, reportMissingModuleSource
+# mypy: disable-error-code="import-untyped"
 
 import logging
 import re
@@ -14,6 +15,7 @@ from datetime import datetime, timezone
 from typing import Optional
 from feedgen.feed import FeedGenerator
 from pyrogram import errors
+from pyrogram.types import Message
 from post_parser import PostParser
 from config import get_settings
 
@@ -33,9 +35,9 @@ async def _create_time_based_media_groups(messages, merge_seconds: int = 5):
     Create media groups based on time difference between messages
     """
     messages_sorted = sorted(messages, key=lambda x: x.date)
-    cluster = []
-    last_msg_date = None
-    current_media_group_id = None
+    cluster: list[Message] = []
+    last_msg_date: Optional[datetime] = None
+    current_media_group_id: Optional[str] = None
 
     for msg in messages_sorted:
         
@@ -51,26 +53,32 @@ async def _create_time_based_media_groups(messages, merge_seconds: int = 5):
         
         if time_diff <= merge_seconds:
             if current_media_group_id:
-                msg.media_group_id = current_media_group_id
+                msg.media_group_id = current_media_group_id  # type: ignore
             elif msg_media_group_id:
                 current_media_group_id = msg_media_group_id
                 for m in cluster:
-                    m.media_group_id = current_media_group_id
+                    m.media_group_id = current_media_group_id  # type: ignore
             cluster.append(msg)
             last_msg_date = msg.date
         else:
             if len(cluster) >= 2 and not current_media_group_id:
-                new_group_id = f"time_{min(m.date for m in cluster)}"
-                for m in cluster:
-                    m.media_group_id = new_group_id
+                dates = [m.date for m in cluster if m.date is not None]
+                if dates:
+                    min_date = min(dates)
+                    new_group_id = f"time_{min_date}"
+                    for m in cluster:
+                        m.media_group_id = new_group_id  # type: ignore
             cluster = [msg]
             last_msg_date = msg.date
             current_media_group_id = msg_media_group_id
     
     if len(cluster) >= 2 and not current_media_group_id:
-        new_group_id = f"time_{min(m.date for m in cluster)}"
-        for m in cluster:
-            m.media_group_id = new_group_id
+        dates = [m.date for m in cluster if m.date is not None]
+        if dates:
+            min_date = min(dates)
+            new_group_id = f"time_{min_date}"
+            for m in cluster:
+                m.media_group_id = new_group_id  # type: ignore
 
     return messages