--- a/script/stream/recorder_tweetstream.py	Thu Jun 23 10:33:32 2011 +0200
+++ b/script/stream/recorder_tweetstream.py	Fri Jun 24 23:41:52 2011 +0200
@@ -3,6 +3,7 @@
 from optparse import OptionParser
 from sqlalchemy.orm import sessionmaker
 from sqlite3 import *
+import datetime
 import logging
 import os
 import socket
@@ -47,12 +48,12 @@
         self.retry_wait = retry_wait
         self._reconnects = 0
         self._error_cb = error_cb
-        super(ReconnectingTweetStream,self).__init__(auth, keywords, url, **kwargs)
+        super(ReconnectingTweetStream, self).__init__(auth, keywords, url, **kwargs)
 
     def next(self):
         while True:
             try:
-                return super(ReconnectingTweetStream,self).next()
+                return super(ReconnectingTweetStream, self).next()
             except tweetstream.ConnectionError, e:
                 logging.debug("connection error :" + str(e))
                 self._reconnects += 1
@@ -74,12 +75,12 @@
     screen_name = ""
     if 'user' in tweet and 'screen_name' in tweet['user']:
         screen_name = tweet['user']['screen_name']
-    logging.info("Process_tweet from %s : %s" % (screen_name,tweet['text']))
+    logging.info("Process_tweet from %s : %s" % (screen_name, tweet['text']))
     logging.debug("Process_tweet :" + repr(tweet))
     processor = utils.TwitterProcessor(tweet, None, session, token_filename)
     processor.process()
 
-def main(username, password, track, session, debug, reconnects, token_filename):
+def main(username, password, track, session, debug, reconnects, token_filename, duration):
 
     #username = username or raw_input('Twitter username: ')
     #password = password or getpass('Twitter password: ')
@@ -95,9 +96,15 @@
         auth = tweetstream.auth.OAuthHandler(consumer_key, consumer_secret, secure=False)
         auth.set_access_token(*(utils.get_oauth_token(token_filename)))
     
+    if duration >= 0:
+        end_ts = datetime.datetime.utcnow() + datetime.timedelta(seconds=duration)
+    
     stream = ReconnectingTweetStream(auth, track_list, reconnects=reconnects)
     try:
-        for tweet in stream:
+        for tweet in stream:            
+            if duration >= 0 and  datetime.datetime.utcnow() >= end_ts:
+                print "Stop recording after %d seconds." % (duration)
+                break
             process_tweet(tweet, session, debug, token_filename)
             session.commit()
     finally:
@@ -105,7 +112,7 @@
         
 def get_options():
     parser = OptionParser()
-    parser.add_option("-f", "--file", dest="filename",  
+    parser.add_option("-f", "--file", dest="filename",
                       help="write tweet to FILE", metavar="FILE", default="enmi2010_twitter.db")
     parser.add_option("-u", "--user", dest="username",
                       help="Twitter user", metavar="USER", default=None)
@@ -119,6 +126,9 @@
                       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")
+    parser.add_option("-d", "--duration", dest="duration",
+                      help="Duration of recording in seconds", metavar="DURATION", default= -1, type='int')
+
 
     utils.set_logging_options(parser)
 
@@ -139,13 +149,13 @@
     if options.new and os.path.exists(options.filename):
         os.remove(options.filename)
     
-    engine, metadata = models.setup_database('sqlite:///'+options.filename, echo=(options.debug))
+    engine, metadata = models.setup_database('sqlite:///' + options.filename, echo=(options.debug))
     Session = sessionmaker(bind=engine)
     session = Session()
 
     try:
         try:
-            main(options.username, options.password, options.track, session, options.debug, options.reconnects, options.token_filename)
+            main(options.username, options.password, options.track, session, options.debug, options.reconnects, options.token_filename, options.duration)
         except KeyboardInterrupt:
             print '\nGoodbye!'
         session.commit()