From 377b8360138cf0a9493b7d31ca4bda4cd0670fd2 Mon Sep 17 00:00:00 2001
From: Josh Roesslein <jroesslein@gmail.com>
Date: Sat, 8 Aug 2009 01:38:08 -0500
Subject: [PATCH] Polish up stream watcher. Add spitzer, follow, and track
 streaming modes to stream object.

---
 stream_watcher.py   | 29 ++++++++++++++++++++++++-----
 tweepy/streaming.py | 37 ++++++++++++++++++++++++++++---------
 2 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/stream_watcher.py b/stream_watcher.py
index b6bce34..b31fd1f 100644
--- a/stream_watcher.py
+++ b/stream_watcher.py
@@ -7,21 +7,40 @@ def callback(status):
 
   print status.text   
 
+# Prompt for login credentials and setup stream object
 username = raw_input('Twitter username: ')
 password = getpass('Twitter password: ')
-stream = tweepy.Stream('spritzer', username, password)
-stream.connect(callback)
+stream = tweepy.Stream(username, password, callback)
 
+# Prompt for mode of streaming and connect
 while True:
+  mode = raw_input('Mode? [spritzer/follow/track] ')
+  if mode == 'spritzer':
+    stream.spritzer()
+    break
+  elif mode == 'follow':
+    follow_list = raw_input('Users to follow (comma separated): ')
+    stream.follow(follow_list)
+    break
+  elif mode == 'track':
+    track_list = raw_input('Keywords to track (comma separated): ')
+    stream.track(track_list)
+    break
+  else:
+    print 'Invalid choice! Try again.'
 
+# Run in a loop until termination
+while True:
   try:
-    if not stream.running:
+    if stream.running is False:
       print 'Stream stopped!'
       break
-
     time.sleep(1)
-
   except KeyboardInterrupt:
     break
 
+# Shutdown connection
 stream.disconnect()
+
+print 'Bye!'
+
diff --git a/tweepy/streaming.py b/tweepy/streaming.py
index 0cd035d..b2c849d 100644
--- a/tweepy/streaming.py
+++ b/tweepy/streaming.py
@@ -8,6 +8,7 @@ from threading import Thread
 from . auth import BasicAuthHandler
 from . parsers import parse_status
 from . api import API
+from . error import TweepError
 
 try:
   import json
@@ -16,19 +17,19 @@ except ImportError:
 
 class Stream(object):
 
-  def __init__(self, method, username, password, host='stream.twitter.com', buffer_size=1500):
-    self.method = method if method[0] == '/' else '/' + method
+  def __init__(self, username, password, callback, host='stream.twitter.com', buffer_size=1500):
     self.host = host
     self.auth = BasicAuthHandler(username, password)
     self.running = False
     self.buffer_size = buffer_size
+    self.callback = callback
 
   def _run(self):
     api = API()
-    conn = httplib.HTTPConnection(self.host)
+    conn = httplib.HTTPConnection(self.host, timeout=5)
     headers = {}
     self.auth.apply_auth(None, None, headers, None)
-    conn.request('POST', self.method + '.json?delimited=length', headers=headers)
+    conn.request('POST', self.url, headers=headers)
     resp = conn.getresponse()
     data = ''
     while self.running:
@@ -37,7 +38,7 @@ class Stream(object):
 
       # read length
       length = ''
-      while True:
+      while resp.isclosed() is False:
         c = resp.read(1)
         if c == '\n':
           break
@@ -62,15 +63,33 @@ class Stream(object):
     conn.close()
     self.running = False
 
-  def connect(self, callback):
+  def spritzer(self):
     if self.running:
       raise TweepError('Stream object already connected!')
-    self.callback = callback
+    self.url = '/spritzer.json?delimited=length'
+    self.running = True
+    Thread(target=self._run).start()
+
+  def follow(self, follow=None):
+    if self.running:
+      raise TweepError('Stream object already connected!')
+    self.url = '/follow.json?delimited=length'
+    if follow:
+      self.url += '&follow=%s' % str(follow).strip('[]').replace(' ', '')
+    self.running = True
+    Thread(target=self._run).start()
+
+  def track(self, track=None):
+    if self.running:
+      raise TweepError('Stream object already connected!')
+    self.url = '/track.json?delimited=length'
+    if track:
+      self.url += '&track=%s' % str(track).strip('[]').replace(' ', '')
     self.running = True
     Thread(target=self._run).start()
     
   def disconnect(self):
-    if not self.running:
-      raise TweepError('Stream object not connected!')
+    if self.running is False:
+      return
     self.running = False
       
-- 
2.25.1