From 9624b45314b46eff513a90240449a9fae50e153a Mon Sep 17 00:00:00 2001
From: Tiger 1 <polunindima2002@gmail.com>
Date: Thu, 6 Jun 2019 21:47:08 +0300
Subject: [PATCH] add Poll, sendPoll, stopPoll

---
 telebot/__init__.py  |  7 +++++++
 telebot/apihelper.py | 20 ++++++++++++++++++
 telebot/types.py     | 48 ++++++++++++++++++++++++++++++++++++++++++++
 tests/my_tests.py    | 41 +++++++++++++++++++++++++++++++++++++
 4 files changed, 116 insertions(+)
 create mode 100644 tests/my_tests.py

diff --git a/telebot/__init__.py b/telebot/__init__.py
index 52a4e5a..98076cd 100644
--- a/telebot/__init__.py
+++ b/telebot/__init__.py
@@ -113,6 +113,7 @@ class TeleBot:
         getUserProfilePhotos
         getUpdates
         getFile
+        sendPoll
         kickChatMember
         unbanChatMember
         restrictChatMember
@@ -1044,6 +1045,12 @@ class TeleBot:
                                         disable_notification, reply_to_message_id, reply_markup, provider_data)
         return types.Message.de_json(result)
 
+    def send_poll(self, chat_id, poll, disable_notifications=False, reply_to_message=None, reply_markup=None):
+        return types.Message.de_json(apihelper.send_poll(self.token, chat_id, poll.question, poll.options, disable_notifications, reply_to_message, reply_markup))
+
+    def stop_poll(self, chat_id, message_id):
+        return types.Poll.de_json(apihelper.stop_poll(self.token, chat_id, message_id))
+
     def answer_shipping_query(self, shipping_query_id, ok, shipping_options=None, error_message=None):
         return apihelper.answer_shipping_query(self.token, shipping_query_id, ok, shipping_options, error_message)
 
diff --git a/telebot/apihelper.py b/telebot/apihelper.py
index 749fc98..094fbbf 100644
--- a/telebot/apihelper.py
+++ b/telebot/apihelper.py
@@ -938,6 +938,26 @@ def delete_sticker_from_set(token, sticker):
     return _make_request(token, method_url, params=payload, method='post')
 
 
+def send_poll(token, chat_id, question, options, disable_notifications=False, reply_to_message_id=None, reply_markup=None):
+    method_url = r'sendPoll'
+    payload = {'chat_id': str(chat_id), 'question': question, 'options': _convert_list_json_serializable(options)}
+    if disable_notifications:
+        payload['disable_notification'] = disable_notifications
+    if reply_to_message_id:
+        payload['reply_to_message_id'] = reply_to_message_id
+    if reply_markup:
+        payload['reply_markup'] = _convert_markup(reply_markup)
+    return _make_request(token, method_url, params=payload)
+
+
+def stop_poll(token, chat_id, message_id, reply_markup=None):
+    method_url = r'stopPoll'
+    payload = {'chat_id': str(chat_id), 'message_id': message_id}
+    if reply_markup:
+        payload['reply_markup'] = _convert_markup(reply_markup)
+    return _make_request(token, method_url, params=payload)
+
+
 def _convert_list_json_serializable(results):
     ret = ''
     for r in results:
diff --git a/telebot/types.py b/telebot/types.py
index 6165c94..0965ebd 100644
--- a/telebot/types.py
+++ b/telebot/types.py
@@ -368,6 +368,9 @@ class Message(JsonDeserializable):
         if 'connected_website' in obj:
             opts['connected_website'] = obj['connected_website']
             content_type = 'connected_website'
+        if 'poll' in obj:
+            opts['poll'] = Poll.de_json(obj['poll'])
+            content_type = 'poll'
         return cls(message_id, from_user, date, chat, content_type, opts, json_string)
 
     @classmethod
@@ -2173,3 +2176,48 @@ class InputMediaDocument(InputMedia):
         if self.thumb:
             ret['thumb'] = self.thumb
         return ret
+
+
+class PollOption(JsonSerializable, JsonDeserializable):
+    @classmethod
+    def de_json(cls, json_type):
+        obj = cls.check_json(json_type)
+        text = obj['text']
+        voter_count = int(obj['voter_count'])
+        option = cls(text)
+        option.voter_count = voter_count
+        return option
+
+    def __init__(self, text):
+        self.text = text
+        self.voter_count = 0
+
+    def to_json(self):
+        return json.dumps(self.text)
+
+
+class Poll(JsonDeserializable):
+    @classmethod
+    def de_json(cls, json_type):
+        obj = cls.check_json(json_type)
+        poll_id = obj['id']
+        question = obj['question']
+        poll = cls(question)
+        options = []
+        for opt in obj['options']:
+            options.append(PollOption.de_json(opt))
+        poll.options = options
+        is_closed = obj['is_closed']
+        poll.id = poll_id
+        poll.is_closed = is_closed
+        return poll
+
+    def __init__(self, question):
+        self.options = []
+        self.question = question
+
+    def add(self, option):
+        if type(option) is PollOption:
+            self.options.append(option)
+        else:
+            self.options.append(PollOption(option))
diff --git a/tests/my_tests.py b/tests/my_tests.py
new file mode 100644
index 0000000..05d4078
--- /dev/null
+++ b/tests/my_tests.py
@@ -0,0 +1,41 @@
+import os
+from telebot import logger, logging, types, TeleBot
+import telebot.apihelper as api
+
+try:
+    TOKEN = os.environ['TOKEN']
+except:
+    logger.error('Not variable \'TOKEN\' in environ')
+    exit(1)
+
+CHAT_ID = -1001405019571
+logger.setLevel(logging.DEBUG)
+
+bot = TeleBot(TOKEN)
+
+
+@bot.message_handler(content_types=['poll'])
+def po(m):
+    logger.debug('Give poll')
+    bot.send_message(m.chat.id, 'Я тоже так умею!')
+    m = bot.send_poll(m.chat.id, m.poll)
+    print(m.chat.id, m.message_id)
+
+
+def test_send_poll():
+    poll = types.Poll('Какой ты сегодня?')
+    poll.add('Добрый')
+    poll.add('Веселый')
+    poll.add('Грустный')
+    kb = types.InlineKeyboardMarkup()
+    kb.add(types.InlineKeyboardButton('1', url='t.me/dr_forse'))
+    result = bot.send_poll(CHAT_ID, poll, reply_to_message=60312, reply_markup=kb)
+    assert result['poll']['question'] == 'Какой ты сегодня?'
+
+
+def test_stop_poll():
+    res = bot.stop_poll(-1001405019571, 60370)
+
+
+test_stop_poll()
+bot.polling(none_stop=True, timeout=600)