streaming.py 7.85 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

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

14
from tweepy.utils import import_simplejson, urlencode_noplus
15
json = import_simplejson()
16

Joshua Roesslein's avatar
Joshua Roesslein committed
17
STREAM_VERSION = '1.1'
18

Josh Roesslein's avatar
Josh Roesslein committed
19

20
21
class StreamListener(object):

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

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

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

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

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

64

65
66
class Stream(object):

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

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

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

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

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

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

134
        if exception:
135
            raise
136

137
    def _data(self, data):
138
        if self.listener.on_data(data) is False:
139
            self.running = False
140

Josh Roesslein's avatar
Josh Roesslein committed
141
    def _read_loop(self, resp):
142

143
        while self.running and not resp.isclosed():
144
145
146

            # Note: keep-alive newlines might be inserted before each length value.
            # read until we get a digit...
147
148
            c = '\n'
            while c == '\n' and self.running and not resp.isclosed():
149
150
151
152
                c = resp.read(1)
            delimited_string = c

            # read rest of delimiter length..
Steve Jones's avatar
Steve Jones committed
153
            d = ''
154
            while d != '\n' and self.running and not resp.isclosed():
Steve Jones's avatar
Steve Jones committed
155
156
                d = resp.read(1)
                delimited_string += d
157
158

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

163
164
        if resp.isclosed():
            self.on_closed(resp)
Josh Roesslein's avatar
Josh Roesslein committed
165

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

173
174
175
176
    def on_closed(self, resp):
        """ Called when the response has been closed by Twitter """
        pass

AlanBell's avatar
AlanBell committed
177
    def userstream(self, count=None, async=False, secure=True):
178
        self.parameters = {'delimited': 'length'}
AlanBell's avatar
AlanBell committed
179
180
        if self.running:
            raise TweepError('Stream object already connected!')
181
        self.url = '/2/user.json?delimited=length'
AlanBell's avatar
AlanBell committed
182
183
        self.host='userstream.twitter.com'
        self._start(async)
184
185

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

194
    def retweet(self, async=False):
195
        self.parameters = {'delimited': 'length'}
196
197
        if self.running:
            raise TweepError('Stream object already connected!')
Joshua Roesslein's avatar
Joshua Roesslein committed
198
        self.url = '/%s/statuses/retweet.json?delimited=length' % STREAM_VERSION
199
        self._start(async)
200

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

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

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