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

Josh Roesslein's avatar
Josh Roesslein committed
11
from tweepy.auth import BasicAuthHandler
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, timeout=5.0, retry_count = None,
72
73
                    retry_time = 10.0, snooze_time = 5.0, buffer_size=1500,
                    headers=None, secure=False):
74
        self.auth = auth
Josh Roesslein's avatar
Josh Roesslein committed
75
76
77
78
79
80
81
82
        self.running = False
        self.timeout = timeout
        self.retry_count = retry_count
        self.retry_time = retry_time
        self.snooze_time = snooze_time
        self.buffer_size = buffer_size
        self.listener = listener
        self.api = API()
83
        self.headers = headers or {}
84
85
86
87
        if secure:
            self.scheme = "https://"
        else:
            self.scheme = "http://"
88
        self.body = None
Josh Roesslein's avatar
Josh Roesslein committed
89
90
91

    def _run(self):
        # setup
92
93
        url = "%s%s%s" % (self.scheme, self.host, self.url)
        self.auth.apply_auth(url, 'GET', self.headers, None)
Josh Roesslein's avatar
Josh Roesslein committed
94
95
96
97

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

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

133
134
135
        if exception:
            raise exception

Josh Roesslein's avatar
Josh Roesslein committed
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
    def _read_loop(self, resp):
        data = ''
        while self.running:
            if resp.isclosed():
                break

            # read length
            length = ''
            while True:
                c = resp.read(1)
                if c == '\n':
                    break
                length += c
            length = length.strip()
            if length.isdigit():
                length = int(length)
            else:
                continue

155
            # read data and pass into listener
Josh Roesslein's avatar
Josh Roesslein committed
156
            data = resp.read(length)
157
158
            if self.listener.on_data(data) is False:
                self.running = False
Josh Roesslein's avatar
Josh Roesslein committed
159

160
161
162
163
164
165
166
167
    def _start(self, async):
        self.running = True
        if async:
            Thread(target=self._run).start()
        else:
            self._run()

    def firehose(self, count=None, async=False):
Josh Roesslein's avatar
Josh Roesslein committed
168
169
170
171
172
        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
173
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
174

175
    def retweet(self, async=False):
176
177
178
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/retweet.json?delimited=length' % STREAM_VERSION
179
        self._start(async)
180

181
    def sample(self, count=None, async=False):
Josh Roesslein's avatar
Josh Roesslein committed
182
183
184
185
186
        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
187
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
188

189
    def filter(self, follow=None, track=None, async=False, locations=None):
190
191
        params = {}
        self.headers['Content-type'] = "application/x-www-form-urlencoded"
Josh Roesslein's avatar
Josh Roesslein committed
192
193
194
195
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/filter.json?delimited=length' % STREAM_VERSION
        if follow:
196
            params['follow'] = ','.join(map(str, follow))
Josh Roesslein's avatar
Josh Roesslein committed
197
        if track:
198
            params['track'] = ','.join(map(str, track))
199
200
        if locations and len(locations) > 0:
            assert len(locations) % 4 == 0
201
            params['locations'] = ','.join(['%.2f' % l for l in locations])
202
        self.body = urllib.urlencode(params)
203
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
204
205

    def disconnect(self):
206
        if self.running is False:
Josh Roesslein's avatar
Josh Roesslein committed
207
208
209
            return
        self.running = False