script/stream/recorder_tweetstream.py
changeset 15 5d552b6a0e55
parent 11 54d7f1486ac4
child 82 210dc265c70f
--- a/script/stream/recorder_tweetstream.py	Tue Jan 18 18:25:18 2011 +0100
+++ b/script/stream/recorder_tweetstream.py	Thu Jan 20 10:44:04 2011 +0100
@@ -8,6 +8,7 @@
 import socket
 import sys
 import tweetstream
+import tweetstream.auth
 socket._fileobject.default_bufsize = 0
 
 
@@ -41,21 +42,22 @@
 
     """
 
-    def __init__(self, user, password, keywords, url="track", reconnects=3, error_cb=None, retry_wait=5, **kwargs):
+    def __init__(self, auth, keywords, url="track", reconnects=3, error_cb=None, retry_wait=5, **kwargs):
         self.max_reconnects = reconnects
         self.retry_wait = retry_wait
         self._reconnects = 0
         self._error_cb = error_cb
-        super(ReconnectingTweetStream,self).__init__(user, password, keywords, url, **kwargs)
+        super(ReconnectingTweetStream,self).__init__(auth, keywords, url, **kwargs)
 
     def next(self):
         while True:
             try:
                 return super(ReconnectingTweetStream,self).next()
             except tweetstream.ConnectionError, e:
+                logging.debug("connection error :" + str(e))
                 self._reconnects += 1
                 if self._reconnects > self.max_reconnects:
-                    raise ConnectionError("Too many retries")
+                    raise tweetstream.ConnectionError("Too many retries")
 
                 # Note: error_cb is not called on the last error since we
                 # raise a ConnectionError instead
@@ -68,24 +70,32 @@
 
 
 
-def process_tweet(tweet, session, debug):
+def process_tweet(tweet, session, debug, token_filename):
     
     logging.debug("Process_tweet :" + repr(tweet))
-    processor = utils.TwitterProcessor(tweet, None, session)
+    processor = utils.TwitterProcessor(tweet, None, session, token_filename)
     processor.process()
 
-def main(username, password, track, session, debug, reconnects):
+def main(username, password, track, session, debug, reconnects, token_filename):
 
-    username = username or raw_input('Twitter username: ')
-    password = password or getpass('Twitter password: ')
+    #username = username or raw_input('Twitter username: ')
+    #password = password or getpass('Twitter password: ')
 
     track_list = track or raw_input('Keywords to track (comma seperated): ').strip()
     track_list = [k for k in track_list.split(',')]
-
-    stream = ReconnectingTweetStream(username, password, track_list, reconnects=reconnects)
+    
+    if username and password:
+        auth = tweetstream.auth.BasicAuthHandler(username, password)        
+    else:
+        consumer_key = models.CONSUMER_KEY
+        consumer_secret = models.CONSUMER_SECRET
+        auth = tweetstream.auth.OAuthHandler(consumer_key, consumer_secret, secure=False)
+        auth.set_access_token(*(utils.get_oauth_token(token_filename)))
+    
+    stream = ReconnectingTweetStream(auth, track_list, reconnects=reconnects)
     try:
         for tweet in stream:
-            process_tweet(tweet, session, debug)
+            process_tweet(tweet, session, debug, token_filename)
             session.commit()
     finally:
         stream.close()
@@ -98,13 +108,15 @@
                       help="Twitter user", metavar="USER", default=None)
     parser.add_option("-w", "--password", dest="password",
                       help="Twitter password", metavar="PASSWORD", default=None)
-    parser.add_option("-t", "--track", dest="track",
+    parser.add_option("-T", "--track", dest="track",
                       help="Twitter track", metavar="TRACK")
     parser.add_option("-n", "--new", dest="new", action="store_true",
                       help="new database", default=False)
     parser.add_option("-r", "--reconnects", dest="reconnects",
                       help="Reconnects", metavar="RECONNECTS", default=10, type='int')
-    
+    parser.add_option("-t", dest="token_filename", metavar="TOKEN_FILENAME", default=".oauth_token",
+                      help="Token file name")
+
     utils.set_logging_options(parser)
 
     return parser.parse_args()
@@ -130,7 +142,7 @@
 
     try:
         try:
-            main(options.username, options.password, options.track, session, options.debug, options.reconnects)
+            main(options.username, options.password, options.track, session, options.debug, options.reconnects, options.token_filename)
         except KeyboardInterrupt:
             print '\nGoodbye!'
         session.commit()