streaming.py 6.19 KB
Newer Older
1
2
3
4
5
# Tweepy
# Copyright 2009 Joshua Roesslein
# See LICENSE

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
12
13
14
from tweepy.auth import BasicAuthHandler
from tweepy.parsers import parse_status
from tweepy.api import API
from tweepy.error import TweepError
15
16

try:
17
    import simplejson as json
18
except ImportError:
19
    try:
20
        import json  # Python 2.6+
21
22
    except ImportError:
        try:
23
            from django.utils import simplejson as json  # Google App Engine
24
25
        except ImportError:
            raise ImportError, "Can't load a json library"
26

27
28
STREAM_VERSION = 1

Josh Roesslein's avatar
Josh Roesslein committed
29

30
31
class StreamListener(object):

32
33
    def __init__(self, api=None):
        self.api = api or API()
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

    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:
            status = parse_status(json.loads(data), self.api)
            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
54
55
56
57
58
59
60
    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
61

Josh Roesslein's avatar
Josh Roesslein committed
62
63
64
    def on_limit(self, track):
        """Called when a limitation notice arrvies"""
        return
65

Josh Roesslein's avatar
Josh Roesslein committed
66
67
68
    def on_error(self, status_code):
        """Called when a non-200 status code is returned"""
        return False
69

Josh Roesslein's avatar
Josh Roesslein committed
70
71
72
    def on_timeout(self):
        """Called when stream connection times out"""
        return
73

74

75
76
class Stream(object):

Josh Roesslein's avatar
Josh Roesslein committed
77
78
79
80
81
82
83
84
85
86
87
88
89
    host = 'stream.twitter.com'

    def __init__(self, username, password, listener, timeout=5.0, retry_count = None,
                    retry_time = 10.0, snooze_time = 5.0, buffer_size=1500):
        self.auth = BasicAuthHandler(username, password)
        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()
90
91
        self.headers = {}
        self.body = None
Josh Roesslein's avatar
Josh Roesslein committed
92
93
94

    def _run(self):
        # setup
95
        self.auth.apply_auth(None, None, self.headers, None)
Josh Roesslein's avatar
Josh Roesslein committed
96
97
98
99
100
101
102
103
104
105
106
107

        # enter loop
        error_counter = 0
        conn = None
        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)
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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
                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)
            except Exception:
                # any other exception is fatal, so kill loop
                break

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

    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

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

158
159
160
161
162
163
164
165
    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
166
167
168
169
170
        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
171
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
172

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

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

187
    def filter(self, follow=None, track=None, async=False):
188
189
        params = {}
        self.headers['Content-type'] = "application/x-www-form-urlencoded"
Josh Roesslein's avatar
Josh Roesslein committed
190
191
192
193
        if self.running:
            raise TweepError('Stream object already connected!')
        self.url = '/%i/statuses/filter.json?delimited=length' % STREAM_VERSION
        if follow:
194
            params['follow'] = ','.join(map(str, follow))
Josh Roesslein's avatar
Josh Roesslein committed
195
        if track:
196
197
            params['track'] = ','.join(map(str, track))
        self.body = urllib.urlencode(params)
198
        self._start(async)
Josh Roesslein's avatar
Josh Roesslein committed
199
200

    def disconnect(self):
201
        if self.running is False:
Josh Roesslein's avatar
Josh Roesslein committed
202
203
204
            return
        self.running = False