From d2a857d169307b5f37b2f74253401ffefff0f0f7 Mon Sep 17 00:00:00 2001 From: vvzvlad Date: Tue, 8 Apr 2025 18:16:46 +0300 Subject: [PATCH] Refactor title generation --- post_parser.py | 128 ++++++++++++++++++++++---------------- tests/test_post_parser.py | 35 ++++++----- 2 files changed, 94 insertions(+), 69 deletions(-) diff --git a/post_parser.py b/post_parser.py index 3284891..87aa6b0 100644 --- a/post_parser.py +++ b/post_parser.py @@ -117,68 +117,90 @@ class PostParser: return "Unknown author" def _generate_title(self, message: Message) -> str: - # Priority 1: Media Type (if present and not just a webpage preview/poll) - if message.media and message.media != MessageMediaType.WEB_PAGE and message.media != MessageMediaType.POLL: - if message.media == MessageMediaType.PHOTO: return "πŸ“· Photo" - elif message.media == MessageMediaType.VIDEO: return "πŸŽ₯ Video" - elif message.media == MessageMediaType.ANIMATION: return "🎞 GIF" - elif message.media == MessageMediaType.AUDIO: return "🎡 Audio" - elif message.media == MessageMediaType.VOICE: return "🎀 Voice" - elif message.media == MessageMediaType.VIDEO_NOTE: return "πŸ“± Video circle" - elif message.media == MessageMediaType.STICKER: return "🎯 Sticker" - elif message.media == MessageMediaType.POLL: return "πŸ“Š Poll" - elif message.media == MessageMediaType.DOCUMENT: - # Distinguish PDF documents - if hasattr(message.document, 'mime_type') and 'pdf' in message.document.mime_type.lower(): - return "πŸ“„ PDF Document" - else: - return "πŸ“Ž Document" + """Generate a title for a message, based on its content.""" - # Priority 2: Service Messages - if getattr(message, "channel_chat_created", False): + if getattr(message, "channel_chat_created", False): #Channel created service message return "✨ Channel created" text = message.text or message.caption or '' text_stripped = text.strip() - - # Priority 3: Text Content (if media didn't determine the title) + if text_stripped: - # Check if the text contains only URL(s) - text_without_urls = re.sub(r'https?://[^\s<>"\']+', '', text_stripped) - if not text_without_urls.strip(): # If only URL(s) are present after removing them - # Check for specific domains like YouTube + # Check if there is only HTML or URL + text_has_tags = bool(re.search(r'<[^>]+?>', text_stripped)) + text_has_urls = bool(re.search(r'https?://[^\s<>"\']+', text_stripped)) + + # Clean text from tags and links + clean_text = html.unescape(text) + clean_text = re.sub(r'<[^>]+?>', '', clean_text) + clean_text = re.sub(r'https?://[^\s<>"\']+', '', clean_text) + clean_text = clean_text.strip() + + # If after cleaning there is no meaningful content + if not clean_text: + # If there was a
or other special tags + if text_has_tags and not text_has_urls: + return "❓ Unknown Post" + + # If there was a YouTube link if re.search(r'(?:youtube\.com|youtu\.be)', text_stripped.lower()): - return "πŸŽ₯ YouTube Link" # More specific title for YouTube links - return "πŸ”— Web link" # Generic title for other links + return "πŸŽ₯ YouTube Link" + + # Otherwise, it's a normal web link + return "πŸ”— Web link" + + # Process URL at the beginning of the title + if text_has_urls and "://" in clean_text: + # Remove part of text after URL + clean_text = clean_text.split("://")[0].strip() + + # Process line breaks - take only the first line + first_line = clean_text.split('\n', 1)[0] + first_line = re.sub(r'\.$', '', first_line) #Remove dot + first_line = first_line.strip() - # Use first line of text if it's not just URL(s) - # Clean the text by removing URLs and HTML tags first - clean_text = re.sub(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', '', text) - clean_text = re.sub('<[^<]+?>', '', clean_text) - # Remove empty lines resulting from cleaning - clean_text = '\\n'.join(line.strip() for line in clean_text.split('\\n') if line.strip()) + # Process long strings + cut_at = 37 + if len(first_line) > cut_at: return f"{first_line[:cut_at]}..." + + return first_line - if clean_text: # If there is meaningful text left - first_line = clean_text.split('\\n', maxsplit=1)[0] - max_length = 100 # Max title length - if len(first_line) <= max_length: - return first_line.strip() # Use the first line if it's short enough - - # Trim longer lines intelligently - trimmed = first_line[:max_length] - last_space = trimmed.rfind(' ') - if last_space != -1: # Trim at the last space found - trimmed = trimmed[:last_space] - - # Append ellipsis if trimmed - return f"{trimmed.strip()}..." if trimmed else "πŸ“ Text Post" # Fallback if trimming results in empty string - - # Priority 4: Fallbacks for specific types if text is empty or was just URL(s) - # This handles cases where only a web_page preview exists without significant text/caption - if message.web_page: return "πŸ”— Web link" - - # Final fallback if no other condition met - return "οΏ½οΏ½ Unknown Post" + # Process media content + if message.media: + # Polls + if message.media == MessageMediaType.POLL: + if hasattr(message, 'poll') and hasattr(message.poll, 'question'): + poll_question = message.poll.question.strip() + if poll_question: + return f"πŸ“Š Poll: {poll_question}" + else: + return "πŸ“Š Poll" + + # PDF documents + if message.media == MessageMediaType.DOCUMENT: + if hasattr(message.document, 'mime_type') and 'pdf' in message.document.mime_type.lower(): + return "πŸ“„ PDF Document" + else: + return "πŸ“Ž Document" + + # Other media types + if message.media == MessageMediaType.PHOTO: return "πŸ“· Photo" + if message.media == MessageMediaType.VIDEO: return "πŸŽ₯ Video" + if message.media == MessageMediaType.ANIMATION: return "🎞 GIF" + if message.media == MessageMediaType.AUDIO: return "🎡 Audio" + if message.media == MessageMediaType.VOICE: return "🎀 Voice" + if message.media == MessageMediaType.VIDEO_NOTE: return "πŸ“± Video circle" + if message.media == MessageMediaType.STICKER: return "🎯 Sticker" + + # Web pages + if message.web_page: + if message.web_page.title: + return f"πŸ”— {message.web_page.title}" + else: + return "πŸ”— Web link" + + # If nothing matches + return "❓ Unknown Post" def _format_forward_info(self, message: Message) -> Union[str, None]: if forward_from_chat := getattr(message, "forward_from_chat", None): diff --git a/tests/test_post_parser.py b/tests/test_post_parser.py index 1875e11..ffeec43 100644 --- a/tests/test_post_parser.py +++ b/tests/test_post_parser.py @@ -22,13 +22,14 @@ class TestPostParserGenerateTitle(unittest.TestCase): def tearDown(self): pass - def _create_mock_message(self, media=None, text=None, caption=None, document_mime_type=None, channel_chat_created=False, web_page=None): + def _create_mock_message(self, media=None, text=None, caption=None, document_mime_type=None, channel_chat_created=False, web_page=None, poll=None): message = MagicMock(spec=Message) message.media = media message.text = text message.caption = caption message.web_page = web_page message.channel_chat_created = channel_chat_created + message.poll = poll # Mock document properties if document exists if document_mime_type: @@ -82,11 +83,11 @@ class TestPostParserGenerateTitle(unittest.TestCase): def test_generate_title_text_only(self): message = self._create_mock_message(text="This is the first line.\nThis is the second line.") - self.assertEqual(self.parser._generate_title(message), "This is the first line.\nThis is the second line.") + self.assertEqual(self.parser._generate_title(message), "This is the first line") #first line and remove dot def test_generate_title_text_with_html(self): message = self._create_mock_message(text="Bold text first line.\nItalic second line.") - self.assertEqual(self.parser._generate_title(message), "Bold text first line.\nItalic second line.") + self.assertEqual(self.parser._generate_title(message), "Bold text first line") #first line and remove tags and dot def test_generate_title_text_with_url_and_text(self): message = self._create_mock_message(text="Check out this link: https://example.com") @@ -106,35 +107,33 @@ class TestPostParserGenerateTitle(unittest.TestCase): def test_generate_title_caption_with_url_and_text(self): message = self._create_mock_message(media=MessageMediaType.PHOTO, caption="Look at this photo! https://example.com/image.jpg") - # Media takes priority - self.assertEqual(self.parser._generate_title(message), "πŸ“· Photo") + self.assertEqual(self.parser._generate_title(message), "Look at this photo!") def test_generate_title_long_text_trimming(self): long_text = "This is a very long line of text that definitely exceeds the maximum length allowed for a title, so it should be trimmed intelligently at the last space before the limit." message = self._create_mock_message(text=long_text) - expected_title = "This is a very long line of text that definitely exceeds the maximum length allowed for a title, so..." + expected_title = "This is a very long line of text that..." #cut at 30 symbols self.assertEqual(self.parser._generate_title(message), expected_title) def test_generate_title_long_text_no_space_trimming(self): long_text = "Thisisaverylonglineoftextthatdefinitelyexceedsthemaximumlengthallowedforatitlesoitshouldbetrimmedatthelimitbecausehasnospaces." message = self._create_mock_message(text=long_text) - expected_title = "Thisisaverylonglineoftextthatdefinitelyexceedsthemaximumlengthallowedforatitlesoitshouldbetrimmedatt..." + expected_title = "Thisisaverylonglineoftextthatdefinite..." #cut at 30 symbols self.assertEqual(self.parser._generate_title(message), expected_title) def test_generate_title_text_with_only_html_and_urls(self): - message = self._create_mock_message(text="Link https://another.link") - self.assertEqual(self.parser._generate_title(message), "Link name https://another.link") + self.assertEqual(self.parser._generate_title(message), "Link name") def test_generate_title_empty_text_after_cleaning(self): message = self._create_mock_message(text="
https://link.com \n ") - # Use variable to capture the actual value - result = self.parser._generate_title(message) - self.assertIn("Unknown Post", result) # Just check if it contains "Unknown Post" + self.assertEqual(self.parser._generate_title(message), "πŸ”— Web link") def test_generate_title_webpage_preview_only(self): web_page_mock = MagicMock() + web_page_mock.title = "Web page title" message = self._create_mock_message(web_page=web_page_mock, text=" ") # Text is whitespace only - self.assertEqual(self.parser._generate_title(message), "πŸ”— Web link") + self.assertEqual(self.parser._generate_title(message), "πŸ”— Web page title") def test_generate_title_webpage_preview_ignored_with_media(self): web_page_mock = MagicMock() @@ -153,9 +152,13 @@ class TestPostParserGenerateTitle(unittest.TestCase): self.assertIn("Unknown Post", result) # Just check if it contains "Unknown Post" def test_generate_title_poll_media_type(self): - # Poll media type should be ignored for title generation, text should be used - message = self._create_mock_message(media=MessageMediaType.POLL, text="What is your favorite color?") - self.assertEqual(self.parser._generate_title(message), "What is your favorite color?") + # Create a mock poll object with question + poll_mock = MagicMock() + poll_mock.question = "Как ΠΏΡ€Π°Π²ΠΈΠ»ΡŒΠ½ΠΎ?" + message = self._create_mock_message(media=MessageMediaType.POLL, poll=poll_mock) + + # Test should check that poll.question is used for title + self.assertEqual(self.parser._generate_title(message), "πŸ“Š Poll: Как ΠΏΡ€Π°Π²ΠΈΠ»ΡŒΠ½ΠΎ?") def test_generate_title_webpage_media_type(self): # Webpage media type should be ignored for title generation, text should be used