Filters for Category and file types (python-telegram-bot#1046)

spontanurlaub · jh0ker · commit 38e3b91a878f · 2018-03-16T23:41:48.000+01:00
* Added extra Filters for File Type and Category

Added extra Filters for File Type and Category, based on the provided Mime-Type

* Added tests for the new Filters

Added Tests for the Category and File Types Filters.

* Fixed Tests

Fixed the Tests from the last commit

* Little Fix

* Revert of unwanted changes

* Update filters.py

* Changed the strings

* Changed file_type to mime_type

* Fixed Tests
diff --git a/telegram/ext/filters.py b/telegram/ext/filters.py
@@ -221,6 +221,79 @@ def filter(self, message):
     class _Document(BaseFilter):
         name = 'Filters.document'
 
+        class category(BaseFilter):
+            """This Filter filters documents by their category in the mime-type attribute
+
+            Note:
+                This Filter only filters by the mime_type of the document,
+                    it doesn't check the validity of the document.
+                The user can manipulate the mime-type of a message and
+                    send media with wrong types that don't fit to this handler.
+
+            Examples:
+                Filters.documents.category('audio/') returnes `True` for all types
+                of audio sent as file, for example 'audio/mpeg' or 'audio/x-wav'
+            """
+
+            def __init__(self, category):
+                """Initialize the category you want to filter
+
+                Args:
+                    category (str, optional): category of the media you want to filter"""
+                self.category = category
+                self.name = "Filters.document.category('{}')".format(self.category)
+
+            def filter(self, message):
+                if message.document:
+                    return message.document.mime_type.startswith(self.category)
+
+        application = category('application/')
+        audio = category('audio/')
+        image = category('image/')
+        video = category('video/')
+        text = category('text/')
+
+        class mime_type(BaseFilter):
+            """This Filter filters documents by their mime-type attribute
+
+            Note:
+                This Filter only filters by the mime_type of the document,
+                    it doesn't check the validity of document.
+                The user can manipulate the mime-type of a message and
+                    send media with wrong types that don't fit to this handler.
+
+            Examples:
+                Filters.documents.mime_type('audio/mpeg') filters all audio in mp3 format.
+            """
+
+            def __init__(self, mimetype):
+                """Initialize the category you want to filter
+
+                Args:
+                    filetype (str, optional): mime_type of the media you want to filter"""
+                self.mimetype = mimetype
+                self.name = "Filters.document.mime_type('{}')".format(self.mimetype)
+
+            def filter(self, message):
+                if message.document:
+                    return message.document.mime_type == self.mimetype
+
+        apk = mime_type('application/vnd.android.package-archive')
+        doc = mime_type('application/msword')
+        docx = mime_type('application/vnd.openxmlformats-officedocument.wordprocessingml.document')
+        exe = mime_type('application/x-ms-dos-executable')
+        gif = mime_type('video/mp4')
+        jpg = mime_type('image/jpeg')
+        mp3 = mime_type('audio/mpeg')
+        pdf = mime_type('application/pdf')
+        py = mime_type('text/x-python')
+        svg = mime_type('image/svg+xml')
+        txt = mime_type('text/plain')
+        targz = mime_type('application/x-compressed-tar')
+        wav = mime_type('audio/x-wav')
+        xml = mime_type('application/xml')
+        zip = mime_type('application/zip')
+
         def filter(self, message):
             return bool(message.document)
 
diff --git a/tests/test_filters.py b/tests/test_filters.py
@@ -20,7 +20,7 @@
 
 import pytest
 
-from telegram import Message, User, Chat, MessageEntity
+from telegram import Message, User, Chat, MessageEntity, Document
 from telegram.ext import Filters, BaseFilter
 import re
 
@@ -86,6 +86,105 @@ def test_filters_document(self, message):
         message.document = 'test'
         assert Filters.document(message)
 
+    def test_filters_document_type(self, message):
+        message.document = Document("file_id", mime_type="application/vnd.android.package-archive")
+        assert Filters.document.apk(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.doc(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "application/msword"
+        assert Filters.document.doc(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.docx(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "application/vnd.openxmlformats-" \
+                                     "officedocument.wordprocessingml.document"
+        assert Filters.document.docx(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.exe(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "application/x-ms-dos-executable"
+        assert Filters.document.exe(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.docx(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "video/mp4"
+        assert Filters.document.gif(message)
+        assert Filters.document.video(message)
+        assert not Filters.document.jpg(message)
+        assert not Filters.document.text(message)
+
+        message.document.mime_type = "image/jpeg"
+        assert Filters.document.jpg(message)
+        assert Filters.document.image(message)
+        assert not Filters.document.mp3(message)
+        assert not Filters.document.video(message)
+
+        message.document.mime_type = "audio/mpeg"
+        assert Filters.document.mp3(message)
+        assert Filters.document.audio(message)
+        assert not Filters.document.pdf(message)
+        assert not Filters.document.image(message)
+
+        message.document.mime_type = "application/pdf"
+        assert Filters.document.pdf(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.py(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "text/x-python"
+        assert Filters.document.py(message)
+        assert Filters.document.text(message)
+        assert not Filters.document.svg(message)
+        assert not Filters.document.application(message)
+
+        message.document.mime_type = "image/svg+xml"
+        assert Filters.document.svg(message)
+        assert Filters.document.image(message)
+        assert not Filters.document.txt(message)
+        assert not Filters.document.video(message)
+
+        message.document.mime_type = "text/plain"
+        assert Filters.document.txt(message)
+        assert Filters.document.text(message)
+        assert not Filters.document.targz(message)
+        assert not Filters.document.application(message)
+
+        message.document.mime_type = "application/x-compressed-tar"
+        assert Filters.document.targz(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.wav(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "audio/x-wav"
+        assert Filters.document.wav(message)
+        assert Filters.document.audio(message)
+        assert not Filters.document.xml(message)
+        assert not Filters.document.image(message)
+
+        message.document.mime_type = "application/xml"
+        assert Filters.document.xml(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.zip(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "application/zip"
+        assert Filters.document.zip(message)
+        assert Filters.document.application(message)
+        assert not Filters.document.apk(message)
+        assert not Filters.document.audio(message)
+
+        message.document.mime_type = "image/x-rgb"
+        assert not Filters.document.category("application/")(message)
+        assert not Filters.document.mime_type("application/x-sh")(message)
+        message.document.mime_type = "application/x-sh"
+        assert Filters.document.category("application/")(message)
+        assert Filters.document.mime_type("application/x-sh")(message)
+
     def test_filters_photo(self, message):
         assert not Filters.photo(message)
         message.photo = 'test'
