diff -r 249867084b9f -r 514e0ee0c68a script/stream/recorder_tweetstream.py --- 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()