diff --git a/telegram/ext/filters.py b/telegram/ext/filters.py index 50951be4b..6a6e4cdfb 100644 --- a/telegram/ext/filters.py +++ b/telegram/ext/filters.py @@ -221,6 +221,79 @@ class Filters(object): 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 index 69831b7cd..74ea941e7 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -20,7 +20,7 @@ import datetime 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 @@ class TestFilters(object): 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'