streaming.py 7.53 KB
Newer Older
1
# Tweepy
2
3
# Copyright 2009-2010 Joshua Roesslein
# See LICENSE for details.
4
5

import httplib
6
from socket import timeout
7
from threading import Thread
8
from time import sleep
9
import urllib
10

11
from tweepy.models import Status
Josh Roesslein's avatar
Josh Roesslein committed
12
13
from tweepy.api import API
from tweepy.error import TweepError
14

15
16
from tweepy.utils import import_simplejson
json = import_simplejson()
17

18
19
STREAM_VERSION = 1

Josh Roesslein's avatar
Josh Roesslein committed
20

21
22
class StreamListener(object):

23
24
    def __init__(self, api=None):
        self.api = api or API()
25
26
27
28
29
30
31
32
33

    def on_data(self, data):
        """Called when raw data is received from connection.

        Override this method if you wish to manually handle
        the stream data. Return False to stop stream and close connection.
        """

        if 'in_reply_to_status_id' in data:
34
            status = Status.parse(self.api, json.loads(data))
35
36
37
38
39
40
41
42
43
44
            if self.on_status(status) is False:
                return False
        elif 'delete' in data:
            delete = json.loads(data)['delete']['status']
            if self.on_delete(delete['id'], delete['user_id']) is False:
                return False
        elif 'limit' in data:
            if self.on_limit(json.loads(data)['limit']['track']) is False:
                return False

Josh Roesslein's avatar
Josh Roesslein committed
45
46
47
48
49
50
51
    def on_status(self, status):
        """Called when a new status arrives"""
        return

    def on_delete(self, status_id, user_id):
        """Called when a delete notice arrives for a status"""
        return
52

Josh Roesslein's avatar
Josh Roesslein committed
53
54
55
    def on_limit(self, track):
        """Called when a limitation notice arrvies"""
        return
56

Josh Roesslein's avatar
Josh Roesslein committed
57
58
59
    def on_error(self, status_code):
        """Called when a non-200 status code is returned"""
        return False
60

Josh Roesslein's avatar
Josh Roesslein committed
61
62
63
    def on_timeout(self):
        """Called when stream connection times out"""
        return
64

65

66
67
class Stream(object):

Josh Roesslein's avatar
Josh Roesslein committed
68
69
    host = 'stream.twitter.com'

70
    def __init__(self, auth, listener, **options):
71
        self.auth = auth
Josh Roesslein's avatar
Josh Roesslein committed
72
        self.listener = listener
73
        self.running = False
74
        self.timeout = options.get("timeout", 300.0)
75
        self.retry_count = options.get("retry_count")
76
77
78
        self.retry_time = options.get("retry_time", 10.0)
        self.snooze_time = options.get("snooze_time",  5.0)
        self.buffer_size = options.get("buffer_size",  1500)
79
        if options.get("secure", True):
80
81
82
83
            self.scheme = "https"
        else:
            self.scheme = "http"

Josh Roesslein's avatar
Josh Roesslein committed
84
        self.api = API()
85
        self.headers = options.get("headers") or {}
86
        self.parameters = None
87
        self.body = None
Josh Roesslein's avatar
Josh Roesslein committed
88
89

    def _run(self):
90
        # Authenticate
91
        url = "%s://%s%s" % (self.scheme, self.host, self.url)
Josh Roesslein's avatar
Josh Roesslein committed
92

93
        # Connect and process the stream
Josh Roesslein's avatar
Josh Roesslein committed
94
95
        error_counter = 0
        conn = None
96
        exception = None
Josh Roesslein's avatar
Josh Roesslein committed
97
        while self.running:
98
            if self.retry_count is not None and error_counter > self.retry_count:
Josh Roesslein's avatar
Josh Roesslein committed
99
100
101
                # quit if error count greater than retry count
                break
            try:
102
103
104
105
                if self.scheme == "http":
                    conn = httplib.HTTPConnection(self.host)
                else:
                    conn = httplib.HTTPSConnection(self.host)
106
                self.auth.apply_auth(url, 'POST', self.headers, self.parameters)
Josh Roesslein's avatar
Josh Roesslein committed
107
108
                conn.connect()
                conn.sock.settimeout(self.timeout)
109
                conn.request('POST', self.url, self.body, headers=self.headers)
Josh Roesslein's avatar
Josh Roesslein committed
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
                resp = conn.getresponse()
                if resp.status != 200:
                    if self.listener.on_error(resp.status) is False:
                        break
                    error_counter += 1
                    sleep(self.retry_time)
                else:
                    error_counter = 0
                    self._read_loop(resp)
            except timeout:
                if self.listener.on_timeout() == False:
                    break
                if self.running is False:
                    break
                conn.close()
                sleep(self.snooze_time)
126
            except Exception, exception:
Josh Roesslein's avatar
Josh Roesslein committed
127
128
129
130
131
132
133
134
                # any other exception is fatal, so kill loop
                break

        # cleanup
        self.running = False
        if conn:
            conn.close()

135
        if exception:
136
            raise
137

138
139
140
141
142
    def _data(self, data):
        for d in [dt for dt in data.split('\n') if dt]:
            if self.listener.on_data(d) is False:
                self.running = False

Josh Roesslein's avatar
Josh Roesslein committed
143
    def _read_loop(self, resp):
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
        buf = ''
        while self.running and not resp.isclosed():
            c = resp.read(self.buffer_size)
            idx = c.rfind('\n')
            if idx > -1:
                # There is an index. Store the tail part for later,
                # and process the head part as messages. We use idx + 1
                # as we dont' actually want to store the newline.
                data = buf + c[:idx]
                buf = c[idx + 1:]
                self._data(data)
            else:
                # No newline found, so we add this to our accumulated
                # buffer
                buf += c
Josh Roesslein's avatar
Josh Roesslein committed
159

160
161
        if resp.isclosed():
            self.on_closed(resp)
Josh Roesslein's avatar
Josh Roesslein committed
162

163
164
165
166
167
168
    def _start(self, async):
        self.running = True
        if async:
            Thread(target=self._run).start()
        else:
            self._run()
Josh Roesslein's avatar
Josh Roesslein committed
169

170
171
172
173
    def on_closed(self, resp):
        """ Called when the response has been closed by Twitter """
        pass

AlanBell's avatar
AlanBell committed
174
175
176
177
178
179
180
181
    def userstream(self, count=None, async=False, secure=True):
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/2/user.json'
        self.host='userstream.twitter.com'
        if count:
            self.url += '&count=%s' % count
        self._start(async)
182
183

    def firehose(self, count=None, async=False):
184
        self.parameters = {'delimited': 'length'}
Josh Roesslein's avatar
Josh Roesslein committed
185
186
187
188
189
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/firehose.json?delimited=length' % STREAM_VERSION
        if count:
            self.url += '&count=%s' % count
190
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
191

192
    def retweet(self, async=False):
193
        self.parameters = {'delimited': 'length'}
194
195
196
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/retweet.json?delimited=length' % STREAM_VERSION
197
        self._start(async)
198

199
    def sample(self, count=None, async=False):
200
        self.parameters = {'delimited': 'length'}
Josh Roesslein's avatar
Josh Roesslein committed
201
202
203
204
205
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/sample.json?delimited=length' % STREAM_VERSION
        if count:
            self.url += '&count=%s' % count
206
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
207

208
    def filter(self, follow=None, track=None, async=False, locations=None, count = None):
209
        self.parameters = {}
210
        self.headers['Content-type'] = "application/x-www-form-urlencoded"
Josh Roesslein's avatar
Josh Roesslein committed
211
212
213
214
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/filter.json?delimited=length' % STREAM_VERSION
        if follow:
215
            self.parameters['follow'] = ','.join(map(str, follow))
Josh Roesslein's avatar
Josh Roesslein committed
216
        if track:
217
            self.parameters['track'] = ','.join(map(str, track))
218
219
        if locations and len(locations) > 0:
            assert len(locations) % 4 == 0
220
            self.parameters['locations'] = ','.join(['%.2f' % l for l in locations])
221
222
        if count:
            self.parameters['count'] = count
223
224
        self.body = urllib.urlencode(self.parameters)
        self.parameters['delimited'] = 'length'
225
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
226
227

    def disconnect(self):
228
        if self.running is False:
Josh Roesslein's avatar
Josh Roesslein committed
229
230
231
            return
        self.running = False