Use imghdr instead re to match image headers. Fixes #37

This commit is contained in:
Leandro Toledo 2015-08-15 15:00:28 -03:00
parent 2cecca8324
commit 5dc1e4cac1
2 changed files with 27 additions and 17 deletions

View file

@ -26,8 +26,8 @@ except ImportError:
from urllib2 import urlopen from urllib2 import urlopen
import mimetypes import mimetypes
import os import os
import re
import sys import sys
import imghdr
from .error import TelegramError from .error import TelegramError
@ -128,20 +128,9 @@ class InputFile(object):
Returns: Returns:
The str mimetype of an image. The str mimetype of an image.
""" """
try: image = imghdr.what(None, stream)
header = stream[:10] if image:
return 'image/%s' % image
if re.match(b'GIF8', header):
return 'image/gif'
if re.match(b'\x89PNG', header):
return 'image/png'
if re.match(b'\xff\xd8\xff\xe0\x00\x10JFIF', header) or \
re.match(b'\xff\xd8\xff\xe1(.*){2}Exif', header):
return 'image/jpeg'
except IndexError as e:
raise TelegramError(str(e))
raise TelegramError({'message': 'Could not parse file content'}) raise TelegramError({'message': 'Could not parse file content'})

View file

@ -79,13 +79,27 @@ class BotTest(unittest.TestCase):
chat_id=12173560) chat_id=12173560)
self.assertEqual('AgADAQADr6cxGzU8LQe6q0dMJD2rHYkP2ykABFymiQqJgjxRGGMAAgI', message.photo[0].file_id) self.assertEqual('AgADAQADr6cxGzU8LQe6q0dMJD2rHYkP2ykABFymiQqJgjxRGGMAAgI', message.photo[0].file_id)
def testSendURLPhoto(self): def testSendJPGURLPhoto(self):
'''Test the telegram.Bot sendPhoto method''' '''Test the telegram.Bot sendPhoto method'''
print('Testing sendPhoto - URL') print('Testing testSendJPGURLPhoto - URL')
message = self._bot.sendPhoto(photo=str('http://dummyimage.com/600x400/000/fff.jpg&text=telegram'), message = self._bot.sendPhoto(photo=str('http://dummyimage.com/600x400/000/fff.jpg&text=telegram'),
chat_id=12173560) chat_id=12173560)
self.assertEqual(822, message.photo[0].file_size) self.assertEqual(822, message.photo[0].file_size)
def testSendPNGURLPhoto(self):
'''Test the telegram.Bot sendPhoto method'''
print('Testing testSendPNGURLPhoto - URL')
message = self._bot.sendPhoto(photo=str('http://dummyimage.com/600x400/000/fff.png&text=telegram'),
chat_id=12173560)
self.assertEqual(684, message.photo[0].file_size)
def testSendGIFURLPhoto(self):
'''Test the telegram.Bot sendPhoto method'''
print('Testing testSendGIFURLPhoto - URL')
message = self._bot.sendPhoto(photo=str('http://dummyimage.com/600x400/000/fff.gif&text=telegram'),
chat_id=12173560)
self.assertEqual(684, message.photo[0].file_size)
def testSendAudio(self): def testSendAudio(self):
'''Test the telegram.Bot sendAudio method''' '''Test the telegram.Bot sendAudio method'''
print('Testing sendAudio - File') print('Testing sendAudio - File')
@ -107,6 +121,13 @@ class BotTest(unittest.TestCase):
chat_id=12173560) chat_id=12173560)
self.assertEqual(12948, message.document.file_size) self.assertEqual(12948, message.document.file_size)
def testSendGIFURLDocument(self):
'''Test the telegram.Bot sendDocument method'''
print('Testing sendDocument - File')
message = self._bot.sendPhoto(photo=str('http://dummyimage.com/600x400/000/fff.gif&text=telegram'),
chat_id=12173560)
self.assertEqual(684, message.photo[0].file_size)
def testResendDocument(self): def testResendDocument(self):
'''Test the telegram.Bot sendDocument method''' '''Test the telegram.Bot sendDocument method'''
print('Testing sendDocument - Resend') print('Testing sendDocument - Resend')