python-telegram-bot · spontanurlaub · Oct 21, 2017 · Oct 21, 2017 · Oct 21, 2017 · Oct 21, 2017
diff --git a/telegram/ext/filters.py b/telegram/ext/filters.py
@@ -17,8 +17,8 @@
 # You should have received a copy of the GNU Lesser Public License
 # along with this program.  If not, see [http://www.gnu.org/licenses/].
 """This module contains the Filters for use with the MessageHandler class."""
-from telegram import Chat
 from future.utils import string_types
+from telegram import Chat
 
 
 class BaseFilter(object):
@@ -192,12 +192,126 @@ 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 file_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.file_type("audio/mpeg") filters all audio in mp3 format.
+            """
+
+            def __init__(self, filetype):
+                """Initialize the category you want to filter
+
+                Args:
+                    filetype (str, optional): mime_type of the media you want to filter"""
+                self.filetype = filetype
+                self.name = "Filters.document.Filetype(\"{}\")".format(self.filetype)
+
+            def filter(self, message):
+                if message.document:
+                    return message.document.mime_type == self.filetype
+
+        apk = file_type("application/vnd.android.package-archive")
+        doc = file_type("application/msword")
+        docx = file_type("application/vnd.openxmlformats-officedocument.wordprocessingml.document")
+        exe = file_type("application/x-ms-dos-executable")
+        gif = file_type("video/mp4")
+        jpg = file_type("image/jpeg")
+        mp3 = file_type("audio/mpeg")
+        pdf = file_type("application/pdf")
+        py = file_type("text/x-python")
+        svg = file_type("image/svg+xml")
+        txt = file_type("text/plain")
+        targz = file_type("application/x-compressed-tar")
+        wav = file_type("audio/x-wav")
+        xml = file_type("application/xml")
+        zip = file_type("application/zip")
+
         def filter(self, message):
             return bool(message.document)
 
     document = _Document()
     """:obj:`Filter`: Messages that contain :class:`telegram.Document`."""
 
+    class file_size(BaseFilter):
+        """This Filter filters all messages with a `file_size` attribute."""
+
+        def __init__(self, min=None, max=None):
+            """Initialize the limits of the file_size in the `__init__` method.
+
+            Note:
+                If no limits are given, the filter returns `True` for every message
+                with file_size attritute
+
+            Args:
+                min (int, optional): Minimum `file_size` of the message media in Byte.
+                max (int, optional): Maximum `file_size` of the message media in Byte.
+            """
+            self.min = min
+            self.max = max
+            self.name = "Filters.file_size(min={0}, max={1})".format(self.min, self.max)
+
+        def filter(self, message):
+            if message.audio:
+                filesize = message.audio.file_size
+            elif message.document:
+                filesize = message.document.file_size
+            elif message.photo:
+                filesize = message.photo.file_size
+            elif message.sticker:
+                filesize = message.sticker.file_size
+            elif message.video:
+                filesize = message.video.file_size
+            elif message.voice:
+                filesize = message.voice.file_size
+            else:
+                return False
+            if self.min is not None:
+                if filesize < self.min:
+                    return False
+            if self.max is not None:
+                if filesize > self.max:
+                    return False
+            return True
+
     class _Photo(BaseFilter):
         name = 'Filters.photo'
 

diff --git a/tests/test_filters.py b/tests/test_filters.py
@@ -22,6 +22,7 @@
 
 from telegram import Message, User, Chat, MessageEntity
 from telegram.ext import Filters, BaseFilter
+from telegram.files.document import Document
 
 
 @pytest.fixture(scope='function')
@@ -69,6 +70,112 @@ def test_filters_document(self, message):
         message.document = 'test'
         assert Filters.document(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.file_type("application/x-sh")(message)
+        message.document.mime_type = "application/x-sh"
+        assert Filters.document.category("application/")(message)
+        assert Filters.document.file_type("application/x-sh")(message)
+
+    def test_filters_filesize(self, message):
+        assert not Filters.file_size()(message)
+        message.document = Document("file_id", file_size=1000)
+        assert Filters.file_size()(message)
+        assert Filters.file_size(min=500, max=1500)(message)
+        assert not Filters.file_size(min=1500)(message)
+        assert not Filters.file_size(max=500)(message)
+
     def test_filters_photo(self, message):
         assert not Filters.photo(message)
         message.photo = 'test'