← Back to branch summary

~ubuntu-ru-irc/+junk/irckit

« back to all changes in this revision

Viewing changes to ubuntuhelp/ubuntubot/plugins/Poll/plugin.py

  • Committer: rmPIC30 at gmail
  • Date: 2010-07-13 09:08:58 UTC
  • Revision ID: rmpic30@gmail.com-20100713090858-w5kkmk093hx38cen
Добавляю бота

Show diffs side-by-side

added added

removed removed

Lines of Context:
ubuntuhelp/ubuntubot/plugins/Poll/plugin.py
 
1
# -*- coding:utf-8 -*-
 
2
###
 
3
# Copyright (c) 2003-2005, Stéphan Kochen
 
4
# All rights reserved.
 
5
#
 
6
# Redistribution and use in source and binary forms, with or without
 
7
# modification, are permitted provided that the following conditions are met:
 
8
#
 
9
#   * Redistributions of source code must retain the above copyright notice,
 
10
#     this list of conditions, and the following disclaimer.
 
11
#   * Redistributions in binary form must reproduce the above copyright notice,
 
12
#     this list of conditions, and the following disclaimer in the
 
13
#     documentation and/or other materials provided with the distribution.
 
14
#   * Neither the name of the author of this software nor the name of
 
15
#     contributors to this software may be used to endorse or promote products
 
16
#     derived from this software without specific prior written consent.
 
17
#
 
18
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 
19
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 
20
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 
21
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 
22
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 
23
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 
24
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 
25
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 
26
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 
27
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 
28
# POSSIBILITY OF SUCH DAMAGE.
 
29
###
 
30
 
 
31
import os
 
32
import time
 
33
 
 
34
import supybot.ircmsgs as ircmsgs
 
35
import supybot.dbi as dbi
 
36
import supybot.conf as conf
 
37
import supybot.ircdb as ircdb
 
38
import supybot.utils as utils
 
39
from supybot.commands import *
 
40
import supybot.plugins as plugins
 
41
import supybot.ircutils as ircutils
 
42
import supybot.callbacks as callbacks
 
43
 
 
44
class PollError(Exception):
 
45
    pass
 
46
 
 
47
class OptionRecord(dbi.Record):
 
48
    __fields__ = [
 
49
        'text',
 
50
        'votes'
 
51
        ]
 
52
    def __str__(self):
 
53
        return format('#%i: %q', self.id, self.text)
 
54
 
 
55
class PollRecord(dbi.Record):
 
56
    __fields__ = [
 
57
        'by',
 
58
        'question',
 
59
        'options',
 
60
        'status'
 
61
        ]
 
62
    def __str__(self):
 
63
        user = plugins.getUserName(self.by)
 
64
        if self.options:
 
65
            options = format('Options: %s', '; '.join(map(str, self.options)))
 
66
        else:
 
67
            options = 'The poll has no options, yet'
 
68
        if self.status:
 
69
            status = 'open'
 
70
        else:
 
71
            status = 'closed'
 
72
        return format('Poll #%i: %q started by %s. %s.  Poll is %s.',
 
73
                      self.id, self.question, user, options, status)
 
74
 
 
75
class SqlitePollDB(object):
 
76
    def __init__(self, filename):
 
77
        self.dbs = ircutils.IrcDict()
 
78
        self.filename = filename
 
79
 
 
80
    def close(self):
 
81
        for db in self.dbs.itervalues():
 
82
            db.close()
 
83
 
 
84
    def _getDb(self, channel):
 
85
        try:
 
86
            import sqlite
 
87
        except ImportError:
 
88
            raise callbacks.Error, 'You need to have PySQLite installed to ' \
 
89
                                   'use Poll.  Download it at ' \
 
90
                                   '<http://pysqlite.org/>'
 
91
        filename = plugins.makeChannelFilename(self.filename, channel)
 
92
        if filename in self.dbs:
 
93
            return self.dbs[filename]
 
94
        if os.path.exists(filename):
 
95
            self.dbs[filename] = sqlite.connect(filename)
 
96
            return self.dbs[filename]
 
97
        db = sqlite.connect(filename)
 
98
        self.dbs[filename] = db
 
99
        cursor = db.cursor()
 
100
        cursor.execute("""CREATE TABLE polls (
 
101
                          id INTEGER PRIMARY KEY,
 
102
                          question TEXT UNIQUE ON CONFLICT IGNORE,
 
103
                          started_by INTEGER,
 
104
                          open INTEGER)""")
 
105
        cursor.execute("""CREATE TABLE options (
 
106
                          id INTEGER,
 
107
                          poll_id INTEGER,
 
108
                          option TEXT,
 
109
                          UNIQUE (poll_id, id) ON CONFLICT IGNORE)""")
 
110
        cursor.execute("""CREATE TABLE votes (
 
111
                          user_id INTEGER,
 
112
                          poll_id INTEGER,
 
113
                          option_id INTEGER,
 
114
                          UNIQUE (user_id, poll_id)
 
115
                          ON CONFLICT IGNORE)""")
 
116
        db.commit()
 
117
        return db
 
118
 
 
119
    def get(self, channel, poll_id):
 
120
        db = self._getDb(channel)
 
121
        cursor = db.cursor()
 
122
        cursor.execute("""SELECT question, started_by, open
 
123
                          FROM polls WHERE id=%s""", poll_id)
 
124
        if cursor.rowcount:
 
125
            (question, by, status) = cursor.fetchone()
 
126
        else:
 
127
            raise dbi.NoRecordError
 
128
        cursor.execute("""SELECT id, option FROM options WHERE poll_id=%s""",
 
129
                       poll_id)
 
130
        if cursor.rowcount:
 
131
            options = [OptionRecord(i, text=o, votes=0)
 
132
                       for (i, o) in cursor.fetchall()]
 
133
        else:
 
134
            options = []
 
135
        return PollRecord(poll_id, question=question, status=status, by=by,
 
136
                          options=options)
 
137
 
 
138
    def open(self, channel, user, question):
 
139
        db = self._getDb(channel)
 
140
        cursor = db.cursor()
 
141
        cursor.execute("""INSERT INTO polls VALUES (NULL, %s, %s, 1)""",
 
142
                       question, user.id)
 
143
        db.commit()
 
144
        cursor.execute("""SELECT id FROM polls WHERE question=%s""", question)
 
145
        return cursor.fetchone()[0]
 
146
 
 
147
    def closePoll(self, channel, id):
 
148
        db = self._getDb(channel)
 
149
        cursor = db.cursor()
 
150
        # Check to make sure that the poll exists
 
151
        cursor.execute("""SELECT id FROM polls WHERE id=%s""", id)
 
152
        if cursor.rowcount == 0:
 
153
            raise dbi.NoRecordError
 
154
        cursor.execute("""UPDATE polls SET open=0 WHERE id=%s""", id)
 
155
        db.commit()
 
156
 
 
157
    def add(self, channel, user, id, option):
 
158
        db = self._getDb(channel)
 
159
        cursor = db.cursor()
 
160
        # Only the poll starter or an admin can add options
 
161
        cursor.execute("""SELECT started_by FROM polls
 
162
                          WHERE id=%s""",
 
163
                          id)
 
164
        if cursor.rowcount == 0:
 
165
            raise dbi.NoRecordError
 
166
        if not ((user.id == cursor.fetchone()[0]) or
 
167
                (ircdb.checkCapability(user.id, 'admin'))):
 
168
            raise PollError, \
 
169
                    'That poll isn\'t yours and you aren\'t an admin.'
 
170
        # and NOBODY can add options once a poll has votes
 
171
        cursor.execute("""SELECT COUNT(user_id) FROM votes
 
172
                          WHERE poll_id=%s""",
 
173
                          id)
 
174
        if int(cursor.fetchone()[0]) != 0:
 
175
            raise PollError, 'Cannot add options to a poll with votes.'
 
176
        # Get the next highest id
 
177
        cursor.execute("""SELECT MAX(id)+1 FROM options
 
178
                          WHERE poll_id=%s""",
 
179
                          id)
 
180
        option_id = cursor.fetchone()[0] or 1
 
181
        cursor.execute("""INSERT INTO options VALUES
 
182
                          (%s, %s, %s)""",
 
183
                          option_id, id, option)
 
184
        db.commit()
 
185
 
 
186
    def vote(self, channel, user, id, option):
 
187
        db = self._getDb(channel)
 
188
        cursor = db.cursor()
 
189
        cursor.execute("""SELECT open
 
190
                          FROM polls WHERE id=%s""",
 
191
                          id)
 
192
        if cursor.rowcount == 0:
 
193
            raise dbi.NoRecordError
 
194
        elif int(cursor.fetchone()[0]) == 0:
 
195
            raise PollError, 'That poll is closed.'
 
196
        cursor.execute("""SELECT id FROM options
 
197
                          WHERE poll_id=%s
 
198
                          AND id=%s""",
 
199
                          id, option)
 
200
        if cursor.rowcount == 0:
 
201
            raise PollError, 'There is no such option.'
 
202
        cursor.execute("""SELECT option_id FROM votes
 
203
                          WHERE user_id=%s AND poll_id=%s""",
 
204
                          user.id, id)
 
205
        if cursor.rowcount == 0:
 
206
            cursor.execute("""INSERT INTO votes VALUES (%s, %s, %s)""",
 
207
                           user.id, id, option)
 
208
        else:
 
209
            cursor.execute("""UPDATE votes SET option_id=%s
 
210
                              WHERE user_id=%s AND poll_id=%s""",
 
211
                              option, user.id, id)
 
212
        db.commit()
 
213
 
 
214
    def results(self, channel, poll_id):
 
215
        db = self._getDb(channel)
 
216
        cursor = db.cursor()
 
217
        cursor.execute("""SELECT id, question, started_by, open
 
218
                          FROM polls WHERE id=%s""",
 
219
                          poll_id)
 
220
        if cursor.rowcount == 0:
 
221
            raise dbi.NoRecordError
 
222
        (id, question, by, status) = cursor.fetchone()
 
223
        by = plugins.getUserName(by)
 
224
        cursor.execute("""SELECT count(user_id), option_id
 
225
                          FROM votes
 
226
                          WHERE poll_id=%s
 
227
                          GROUP BY option_id
 
228
                          UNION
 
229
                          SELECT 0, id AS option_id
 
230
                          FROM options
 
231
                          WHERE poll_id=%s
 
232
                          AND id NOT IN (
 
233
                                SELECT option_id FROM votes
 
234
                                WHERE poll_id=%s)
 
235
                          GROUP BY option_id
 
236
                          ORDER BY count(user_id) DESC""",
 
237
                          poll_id, poll_id, poll_id)
 
238
        if cursor.rowcount == 0:
 
239
            raise PollError, 'This poll has no votes yet.'
 
240
        else:
 
241
            options = []
 
242
            for count, option_id in cursor.fetchall():
 
243
                cursor.execute("""SELECT option FROM options
 
244
                                  WHERE id=%s AND poll_id=%s""",
 
245
                                  option_id, poll_id)
 
246
                option = cursor.fetchone()[0]
 
247
                options.append(OptionRecord(option_id, votes=int(count),
 
248
                                            text=option))
 
249
        return PollRecord(poll_id, question=question, status=status, by=by,
 
250
                          options=options)
 
251
 
 
252
    def select(self, channel):
 
253
        db = self._getDb(channel)
 
254
        cursor = db.cursor()
 
255
        cursor.execute("""SELECT id, started_by, question
 
256
                          FROM polls
 
257
                          WHERE open=1""")
 
258
        if cursor.rowcount:
 
259
            return [PollRecord(id, question=q, by=by, status=1)
 
260
                    for (id, by, q) in cursor.fetchall()]
 
261
        else:
 
262
            raise dbi.NoRecordError
 
263
 
 
264
PollDB = plugins.DB('Poll', {'sqlite': SqlitePollDB})
 
265
 
 
266
class Poll(callbacks.Plugin):
 
267
    def __init__(self, irc):
 
268
        self.__parent = super(Poll, self)
 
269
        self.__parent.__init__(irc)
 
270
        self.db = PollDB()
 
271
 
 
272
    def die(self):
 
273
        self.__parent.die()
 
274
        self.db.close()
 
275
 
 
276
    def poll(self, irc, msg, args, channel, id):
 
277
        """[<channel>] <id>
 
278
 
 
279
        Displays the poll question and options for the given poll id.
 
280
        <channel> is only necessary if the message isn't sent in the channel
 
281
        itself.
 
282
        """
 
283
        try:
 
284
            record = self.db.get(channel, id)
 
285
        except dbi.NoRecordError:
 
286
            irc.error(format('There is no poll with id %i.', id), Raise=True)
 
287
        irc.reply(record)
 
288
    poll = wrap(poll, ['channeldb', 'id'])
 
289
 
 
290
    def open(self, irc, msg, args, channel, user, question):
 
291
        """[<channel>] <question>
 
292
 
 
293
        Creates a new poll with the given question.  <channel> is only
 
294
        necessary if the message isn't sent in the channel itself.
 
295
        """
 
296
        irc.replySuccess(format('(poll #%i added)',
 
297
                                self.db.open(channel, user, question)))
 
298
    open = wrap(open, ['channeldb', 'user', 'text'])
 
299
 
 
300
    def close(self, irc, msg, args, channel, id):
 
301
        """[<channel>] <id>
 
302
 
 
303
        Closes the poll with the given <id>; further votes will not be allowed.
 
304
        <channel> is only necessary if the message isn't sent in the channel
 
305
        itself.
 
306
        """
 
307
        try:
 
308
            self.db.closePoll(channel, id)
 
309
            irc.replySuccess()
 
310
        except dbi.NoRecordError:
 
311
            irc.errorInvalid('poll id')
 
312
    close = wrap(close, ['channeldb', ('id', 'poll')])
 
313
 
 
314
    def add(self, irc, msg, args, channel, user, id, option):
 
315
        """[<channel>] <id> <option text>
 
316
 
 
317
        Add an option with the given text to the poll with the given id.
 
318
        <channel> is only necessary if the message isn't sent in the channel
 
319
        itself.
 
320
        """
 
321
        try:
 
322
            self.db.add(channel, user, id, option)
 
323
            irc.replySuccess()
 
324
        except dbi.NoRecordError:
 
325
            irc.errorInvalid('poll id')
 
326
        except PollError, e:
 
327
            irc.error(str(e))
 
328
    add = wrap(add, ['channeldb', 'user', ('id', 'poll'), 'text'])
 
329
 
 
330
    def vote(self, irc, msg, args, channel, user, id, option):
 
331
        """[<channel>] <poll id> <option id>
 
332
 
 
333
        Vote for the option with the given id on the poll with the given poll
 
334
        id.  This command can also be used to override any previous vote.
 
335
        <channel> is only necesssary if the message isn't sent in the channel
 
336
        itself.
 
337
        """
 
338
        try:
 
339
            self.db.vote(channel, user, id, option)
 
340
            irc.replySuccess()
 
341
        except dbi.NoRecordError:
 
342
            irc.errorInvalid('poll id')
 
343
        except PollError, e:
 
344
            irc.error(str(e))
 
345
    vote = wrap(vote, ['channeldb', 'user', ('id', 'poll'), ('id', 'option')])
 
346
 
 
347
    def results(self, irc, msg, args, channel, id):
 
348
        """: Используйте !vote в привате, для получения информации по голосованию """
 
349
        try:
 
350
            poll = self.db.results(channel, id)
 
351
            reply = format('Результаты для голосования #%i: %q by %s',
 
352
                           poll.id, poll.question, poll.by)
 
353
            options = poll.options
 
354
            L = []
 
355
            for option in options:
 
356
                L.append(format('#%i %q = %i', option.id, option.text, option.votes))
 
357
            reply += format(' - %L', L)
 
358
            #irc.reply(reply)
 
359
            irc.queueMsg(ircmsgs.privmsg(msg.nick, reply))
 
360
        except dbi.NoRecordError:
 
361
            irc.error('There is no such poll.', Raise=True)
 
362
        except PollError, e:
 
363
            irc.error(str(e))
 
364
    results = wrap(results, ['channeldb', ('id', 'poll')])
 
365
 
 
366
    def list(self, irc, msg, args, channel):
 
367
        """[<channel>]
 
368
 
 
369
        Lists the currently open polls for <channel>.  <channel> is only
 
370
        necessary if the message isn't sent in the channel itself.
 
371
        """
 
372
        try:
 
373
            polls = self.db.select(channel)
 
374
            polls = [format('#%i: %q', p.id, p.question)
 
375
                     for p in polls]
 
376
            irc.reply(format('%L', polls))
 
377
        except dbi.NoRecordError:
 
378
            irc.reply('This channel currently has no open polls.')
 
379
    list = wrap(list, ['channeldb'])
 
380
 
 
381
 
 
382
Class = Poll
 
383
 
 
384
 
 
385
# vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79: