streaming.py 7.66 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
from string import digits
11

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

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

19
20
STREAM_VERSION = 1

Josh Roesslein's avatar
Josh Roesslein committed
21

22
23
class StreamListener(object):

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

    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:
35
            status = Status.parse(self.api, json.loads(data))
36
37
38
39
40
41
42
43
44
45
            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
46
47
48
49
50
51
52
    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
53

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

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

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

66

67
68
class Stream(object):

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

71
    def __init__(self, auth, listener, **options):
72
        self.auth = auth
Josh Roesslein's avatar
Josh Roesslein committed
73
        self.listener = listener
74
        self.running = False
75
        self.timeout = options.get("timeout", 300.0)
76
        self.retry_count = options.get("retry_count")
77
78
79
        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)
80
        if options.get("secure", True):
81
82
83
84
            self.scheme = "https"
        else:
            self.scheme = "http"

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

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

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

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

136
        if exception:
137
            raise
138

139
140
141
142
143
    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
144
    def _read_loop(self, resp):
145

146
        while self.running and not resp.isclosed():
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

            # Note: keep-alive newlines might be inserted before each length value.
            # read until we get a digit...
            c = ''
            while c not in digits:
                c = resp.read(1)
            delimited_string = c

            # read rest of delimiter length..
            while 1:
                d = resp.read(1)
                if d in digits:
                    delimited_string += d
                else:
                    break

            # read the next twitter status object
            if delimited_string.isdigit():
                next_status_obj = resp.read( int(delimited_string) )
                self._data(next_status_obj)
Josh Roesslein's avatar
Josh Roesslein committed
167

168
169
        if resp.isclosed():
            self.on_closed(resp)
Josh Roesslein's avatar
Josh Roesslein committed
170

171
172
173
174
175
176
    def _start(self, async):
        self.running = True
        if async:
            Thread(target=self._run).start()
        else:
            self._run()
Josh Roesslein's avatar
Josh Roesslein committed
177

178
179
180
181
    def on_closed(self, resp):
        """ Called when the response has been closed by Twitter """
        pass

AlanBell's avatar
AlanBell committed
182
183
184
185
186
187
188
189
    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)
190
191

    def firehose(self, count=None, async=False):
192
        self.parameters = {'delimited': 'length'}
Josh Roesslein's avatar
Josh Roesslein committed
193
194
195
196
197
        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
198
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
199

200
    def retweet(self, async=False):
201
        self.parameters = {'delimited': 'length'}
202
203
204
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/retweet.json?delimited=length' % STREAM_VERSION
205
        self._start(async)
206

207
    def sample(self, count=None, async=False):
208
        self.parameters = {'delimited': 'length'}
Josh Roesslein's avatar
Josh Roesslein committed
209
210
211
212
213
        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
214
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
215

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

    def disconnect(self):
236
        if self.running is False:
Josh Roesslein's avatar
Josh Roesslein committed
237
238
239
            return
        self.running = False