<reponame>mihaivavram/twitterbots
from __future__ import absolute_import, print_function

from tweepy.streaming import StreamListener
from tweepy import OAuthHandler, Stream, API
from queue import Queue

from collection.tweets import extract_users

import os
import json
import logging
import time

logger = logging.getLogger(__name__)


def start_lookup(api, lookup_queue, account_queue):
    """Iterates through results in a lookup queue, fetching and returning the
    hydrated account objects.

    This service looks up user ID's 100 ID's at a time per Twitter's API.

    Arguments:
        lookup_queue {queue.Queue} -- The queue to receive lookup requests on
        account_queue {queue.Queue} -- The queue to send hydrated accounts
    """
    logger.info('Account lookup service started')
    try:
        account_batch = []
        # This is a mirror of account_batch but includes the original
        # source so we can tell where an account came from.
        account_sources = {}
        while True:
            account = lookup_queue.get()
            account_id = account['id_str']
            source = account['_tbsource']
            # This probably won't happen, but it'll save us a request anyway
            if account_id in account_sources:
                continue
            account_batch.append(account_id)
            account_sources[account_id] = source
            if len(account_batch) >= 100:
                try:
                    results = api.lookup_users(
                        user_ids=account_batch[:100], include_entities=True)
                except Exception as e:
                    # If an exception occurs, we should probably cool off for
                    # a few seconds and then go ahead and prune out the
                    # results.
                    logger.error(e)
                    logger.info(
                        'Error fetching users... Sleeping for 10 seconds')
                    time.sleep(10)
                    account_batch = account_batch[min(len(account_batch),
                                                      100):]
                    continue
                for result in results:
                    user = result._json
                    # Backfill the original source
                    user['_tbsource'] = account_sources.pop(user['id_str'])
                    account_queue.put(user)
                account_batch = account_batch[min(len(account_batch), 100):]
    except KeyboardInterrupt:
        return


class StdOutListener(StreamListener):
    """ A listener handles tweets that are received from the stream.
    This is a basic listener that just prints received tweets to stdout.
    """

    def __init__(self, api, account_queue, lookup_queue, *args, **kwargs):
        self.account_queue = account_queue
        self.lookup_queue = lookup_queue
        self.source = 'stream'
        super(StdOutListener, self).__init__(*args, api=api, **kwargs)

    def on_data(self, tweet):
        tweet = json.loads(tweet)

        # In the future, we can handle the cases where the message is to
        # delete a tweet. But for now we'll short-circuit
        if 'user' not in tweet:
            return True

        # Handle the user of the tweet. We can just save this off since it
        # contains the full user object
        user = tweet['user']
        user['_tbsource'] = self.source
        self.account_queue.put(user)

        # Extract any additional user IDs that might be in the tweet object
        for account_id in extract_users(tweet):
            self.lookup_queue.put({
                'id_str': account_id,
                '_tbsource': self.source
            })
        return True

    def on_error(self, status):
        print(status)


def start_streamer(api, account_queue, lookup_queue, query=None):
    """Starts a Tweet stream to find account information.

    Arguments:
        api {tweepy.API} -- The authenticated API instance.
    """
    l = StdOutListener(api, account_queue, lookup_queue)
    stream = Stream(api.auth, l)
    if query:
        stream.filter(track=[query])
    else:
        stream.sample()


if __name__ == '__main__':
    consumer_key = os.environ.get('TWEEPY_CONSUMER_KEY')
    consumer_secret = os.environ.get('TWEEPY_CONSUMER_SECRET')
    access_token = os.environ.get('TWEEPY_ACCESS_TOKEN')
    access_token_secret = os.environ.get('TWEEPY_ACCESS_TOKEN_SECRET')

    auth = OAuthHandler(consumer_key, consumer_secret)
    auth.set_access_token(access_token, access_token_secret)

    api = API(auth, wait_on_rate_limit_notify=True, wait_on_rate_limit=True)
    account_queue = Queue()
    lookup_queue = Queue()
    try:
        start_streamer(api, account_queue, lookup_queue)
    except KeyboardInterrupt:
        print('CTRL+C received... shutting down')

from __future__ import absolute_import, print_function

from tweepy.streaming import StreamListener
from tweepy import OAuthHandler, Stream, API
from queue import Queue

from collection.tweets import extract_users

import os
import json
import logging
import time

logger = logging.getLogger(__name__)


def start_lookup(api, lookup_queue, account_queue):
    """Iterates through results in a lookup queue, fetching and returning the
    hydrated account objects.

    This service looks up user ID's 100 ID's at a time per Twitter's API.

    Arguments:
        lookup_queue {queue.Queue} -- The queue to receive lookup requests on
        account_queue {queue.Queue} -- The queue to send hydrated accounts
    """
    logger.info('Account lookup service started')
    try:
        account_batch = []
        # This is a mirror of account_batch but includes the original
        # source so we can tell where an account came from.
        account_sources = {}
        while True:
            account = lookup_queue.get()
            account_id = account['id_str']
            source = account['_tbsource']
            # This probably won't happen, but it'll save us a request anyway
            if account_id in account_sources:
                continue
            account_batch.append(account_id)
            account_sources[account_id] = source
            if len(account_batch) >= 100:
                try:
                    results = api.lookup_users(
                        user_ids=account_batch[:100], include_entities=True)
                except Exception as e:
                    # If an exception occurs, we should probably cool off for
                    # a few seconds and then go ahead and prune out the
                    # results.
                    logger.error(e)
                    logger.info(
                        'Error fetching users... Sleeping for 10 seconds')
                    time.sleep(10)
                    account_batch = account_batch[min(len(account_batch),
                                                      100):]
                    continue
                for result in results:
                    user = result._json
                    # Backfill the original source
                    user['_tbsource'] = account_sources.pop(user['id_str'])
                    account_queue.put(user)
                account_batch = account_batch[min(len(account_batch), 100):]
    except KeyboardInterrupt:
        return


class StdOutListener(StreamListener):
    """ A listener handles tweets that are received from the stream.
    This is a basic listener that just prints received tweets to stdout.
    """

    def __init__(self, api, account_queue, lookup_queue, *args, **kwargs):
        self.account_queue = account_queue
        self.lookup_queue = lookup_queue
        self.source = 'stream'
        super(StdOutListener, self).__init__(*args, api=api, **kwargs)

    def on_data(self, tweet):
        tweet = json.loads(tweet)

        # In the future, we can handle the cases where the message is to
        # delete a tweet. But for now we'll short-circuit
        if 'user' not in tweet:
            return True

        # Handle the user of the tweet. We can just save this off since it
        # contains the full user object
        user = tweet['user']
        user['_tbsource'] = self.source
        self.account_queue.put(user)

        # Extract any additional user IDs that might be in the tweet object
        for account_id in extract_users(tweet):
            self.lookup_queue.put({
                'id_str': account_id,
                '_tbsource': self.source
            })
        return True

    def on_error(self, status):
        print(status)


def start_streamer(api, account_queue, lookup_queue, query=None):
    """Starts a Tweet stream to find account information.

    Arguments:
        api {tweepy.API} -- The authenticated API instance.
    """
    l = StdOutListener(api, account_queue, lookup_queue)
    stream = Stream(api.auth, l)
    if query:
        stream.filter(track=[query])
    else:
        stream.sample()


if __name__ == '__main__':
    consumer_key = os.environ.get('TWEEPY_CONSUMER_KEY')
    consumer_secret = os.environ.get('TWEEPY_CONSUMER_SECRET')
    access_token = os.environ.get('TWEEPY_ACCESS_TOKEN')
    access_token_secret = os.environ.get('TWEEPY_ACCESS_TOKEN_SECRET')

    auth = OAuthHandler(consumer_key, consumer_secret)
    auth.set_access_token(access_token, access_token_secret)

    api = API(auth, wait_on_rate_limit_notify=True, wait_on_rate_limit=True)
    account_queue = Queue()
    lookup_queue = Queue()
    try:
        start_streamer(api, account_queue, lookup_queue)
    except KeyboardInterrupt:
        print('CTRL+C received... shutting down')
Iterates through results in a lookup queue, fetching and returning the     hydrated account objects.      This service looks up user ID's 100 ID's at a time per Twitter's API.      Arguments:         lookup_queue {queue.Queue} -- The queue to receive lookup requests on         account_queue {queue.Queue} -- The queue to send hydrated accounts # This is a mirror of account_batch but includes the original # source so we can tell where an account came from. # This probably won't happen, but it'll save us a request anyway # If an exception occurs, we should probably cool off for # a few seconds and then go ahead and prune out the # results. # Backfill the original source A listener handles tweets that are received from the stream.     This is a basic listener that just prints received tweets to stdout. # In the future, we can handle the cases where the message is to # delete a tweet. But for now we'll short-circuit # Handle the user of the tweet. We can just save this off since it # contains the full user object # Extract any additional user IDs that might be in the tweet object Starts a Tweet stream to find account information.      Arguments:         api {tweepy.API} -- The authenticated API instance.