Added Multi-threaded search manager

GetOldTweets3/manager/ConcurrentTweetManager.py

@@ -0,0 +1,86 @@
+# Author: Giulio Neusch-Frediani - www.github.com/giulionf
+from datetime import datetime
+from datetime import timedelta
+from threading import Thread
+import copy
+
+from GetOldTweets3.manager.TweetManager import TweetManager
+
+DATE_FORMAT = "%Y-%m-%d"
+FIRST_TWEET_DATE = datetime.strptime("2006-03-21", DATE_FORMAT)
+
+
+class ConcurrentTweetManager:
+
+    @staticmethod
+    def getTweets(tweetCriteria, receiveBuffer=None, bufferLength=100, proxy=None, debug=False, worker_count=1,
+                  forceMaxTweets=False):
+
+        if worker_count < 1:
+            raise ValueError("At least one worker is needed")
+
+        if tweetCriteria.maxTweets != 0 and not forceMaxTweets:
+            raise ValueError("Max Tweets is not supported by parallel downloading, since the results can not be ordered"
+                             " by time. If you do not care, you can set forceMaxTweets=True!")
+
+        # Init the queues
+        time_spans = []
+        tweets = []
+        workers = []
+
+        # Split the date in smaller parts
+        since_date = datetime.strptime(tweetCriteria.since, DATE_FORMAT) if hasattr(tweetCriteria, "since") else FIRST_TWEET_DATE
+        until_date = datetime.strptime(tweetCriteria.until, DATE_FORMAT) if hasattr(tweetCriteria, "until") else datetime.now()
+        date_diff_per_worker = (until_date - since_date) / worker_count
+
+        if date_diff_per_worker < timedelta(days=1):
+            max_workers = int((until_date - since_date) / timedelta(days=1))
+            raise ValueError("Too many workers for the time span. Each worker needs at least one day for himself, or"
+                             " some workers will have the same results leading to inconsistencies."
+                             "For your case, the max worker count is {} workers".format(max_workers))
+
+        # Create a TweetCriteria that can be cloned as Model by the Workers
+        criteria = copy.deepcopy(tweetCriteria)
+        if forceMaxTweets:
+            criteria.setMaxTweets(tweetCriteria.maxTweets / worker_count)
+
+        # Create a time span for each of the splitted parts and a corresponding worker
+        for i in range(1, worker_count+1):
+            from_time = since_date + (i-1) * date_diff_per_worker
+            to_time = since_date + i * date_diff_per_worker
+            time_spans.append((from_time, to_time))
+
+        for i in range(0, worker_count):
+            w = WorkerThread(copy.deepcopy(criteria), time_spans[i], tweets, receiveBuffer, bufferLength,
+                             proxy, debug)
+            w.start()
+            workers.append(w)
+
+        # Wait for the workers to finish, then return the results
+        for worker in workers:
+            worker.join()
+
+        return sorted(tweets, key=lambda r: r.date)
+
+
+class WorkerThread(Thread):
+
+    def __init__(self, tweetCriteria, time_span, tweets,  receiveBuffer=None, bufferLength=100, proxy=None,
+                 debug=False):
+        super().__init__()
+        self.stopped = False
+        self.manager = TweetManager()
+        self.tweetCriteria = tweetCriteria
+        self.receiveBuffer = receiveBuffer
+        self.bufferLength = bufferLength
+        self.proxy = proxy
+        self.debug = debug
+        self.time_span = time_span
+        self.tweets = tweets
+
+    def run(self) -> None:
+        self.tweetCriteria.setSince(datetime.strftime(self.time_span[0], "%Y-%m-%d"))
+        self.tweetCriteria.setUntil(datetime.strftime(self.time_span[1], "%Y-%m-%d"))
+        search_results = self.manager.getTweets(self.tweetCriteria, self.receiveBuffer, self.bufferLength, self.proxy,
+                                                self.debug)
+        self.tweets.extend(search_results)

GetOldTweets3/manager/__init__.py

@@ -1,2 +1,3 @@
 from .TweetCriteria import TweetCriteria
-from .TweetManager import TweetManager
+from .TweetManager import TweetManager
+from .ConcurrentTweetManager import ConcurrentTweetManager

bin/GetOldTweets3

@@ -49,6 +49,24 @@ if sys.version_info[0] < 3:
 
 import GetOldTweets3 as got
 
+
+def concatenate_concurrent_files(collected_directory, outputFileName):
+    from os import listdir
+    from os.path import isfile, join
+    file_paths = [collected_directory + "/" + f for f in listdir(collected_directory) if isfile(join(collected_directory, f))]
+
+    with open(outputFileName, "w+", encoding="utf8") as total_result_file:
+        total_result_file.write('date,username,to,replies,retweets,favorites,text,geo,mentions,hashtags,id,permalink\n')
+        for file_path in reversed(file_paths):
+            with open(file_path, "r", encoding="utf8") as file:
+                data = file.read().split("\n", 1)[1]
+                total_result_file.write(data)
+            os.remove(file_path)
+    os.removedirs(collected_directory)
+
+
+
+
 def main(argv):
     if len(argv) == 0:
         print('You must pass some parameters. Use \"-h\" to help.')
@@ -70,13 +88,20 @@ def main(argv):
                                               "maxtweets=",
                                               "lang=",
                                               "output=",
+                                              "debug",
+                                              "workers=",
+                                              "forcemaxtweets",
                                               "emoji=",
                                               "debug"))
 
         tweetCriteria = got.manager.TweetCriteria()
         outputFileName = "output_got.csv"
+        outputDirNameConcurrent = "output_got.csv_collected"
+        outputFileNameConcurrent = "{}_to_{}.csv"
 
         debug = False
+        force_max_tweets = False
+        workers = 1
         usernames = set()
         username_files = set()
         for opt, arg in opts:
@@ -145,10 +170,17 @@ def main(argv):
 
             elif opt == '--output':
                 outputFileName = arg
+                outputDirNameConcurrent = arg + "_collected"
 
             elif opt == '--debug':
                 debug = True
 
+            elif opt == '--workers':
+                workers = int(arg)
+
+            elif opt == '--forcemaxtweets':
+                force_max_tweets = True
+
         if debug:
             print(' '.join(sys.argv))
             print("GetOldTweets3", got.__version__)
@@ -174,8 +206,12 @@ def main(argv):
             else:
                 tweetCriteria.username = usernames.pop()
 
-        outputFile = open(outputFileName, "w+", encoding="utf8")
-        outputFile.write('date,username,to,replies,retweets,favorites,text,geo,mentions,hashtags,id,permalink\n')
+        if workers > 1:
+            if not os.path.exists(outputDirNameConcurrent):
+                os.makedirs(outputDirNameConcurrent)
+        else:
+            outputFile = open(outputFileName, "w+", encoding="utf8")
+            outputFile.write('date,username,to,replies,retweets,favorites,text,geo,mentions,hashtags,id,permalink\n')
 
         cnt = 0
         def receiveBuffer(tweets):
@@ -205,8 +241,45 @@ def main(argv):
             else:
                 print(cnt, end=' ', flush=True)
 
+        def receiveBufferConcurrent(tweets):
+            nonlocal cnt
+
+            last_date = tweets[0].date.strftime("%Y-%m-%d")
+            first_date = tweets[len(tweets)-1].date.strftime("%Y-%m-%d")
+
+            file = open(outputDirNameConcurrent + "/" + outputFileNameConcurrent.format(first_date, last_date), "w+", encoding="utf8")
+            file.write('date,username,to,replies,retweets,favorites,text,geo,mentions,hashtags,id,permalink\n')
+
+            for t in tweets:
+                data_ = [t.date.strftime("%Y-%m-%d %H:%M:%S"),
+                         t.username,
+                         t.to or '',
+                         t.replies,
+                         t.retweets,
+                         t.favorites,
+                         '"' + t.text.replace('"', '""') + '"',
+                         t.geo,
+                         t.mentions,
+                         t.hashtags,
+                         t.id,
+                         t.permalink]
+                data_[:] = [i if isinstance(i, str) else str(i) for i in data_]
+                file.write(','.join(data_) + '\n')
+
+            file.flush()
+            file.close()
+            cnt += len(tweets)
+
+            if sys.stdout.isatty():
+                print("\rSaved %i" % cnt, end='', flush=True)
+            else:
+                print(cnt, end=' ', flush=True)
+
         print("Downloading tweets...")
-        got.manager.TweetManager.getTweets(tweetCriteria, receiveBuffer, debug=debug)
+        if workers > 1:
+            got.manager.ConcurrentTweetManager.getTweets(tweetCriteria, receiveBufferConcurrent, debug=debug, worker_count=workers, forceMaxTweets=force_max_tweets)
+        else:
+            got.manager.TweetManager.getTweets(tweetCriteria, receiveBuffer, debug=debug)
 
     except getopt.GetoptError as err:
         print('Arguments parser error, try -h')
@@ -220,10 +293,14 @@ def main(argv):
         print(str(err))
 
     finally:
-        if "outputFile" in locals():
-            outputFile.close()
-            print()
-            print('Done. Output file generated "%s".' % outputFileName)
+        if "workers" in locals() and workers > 1:
+            concatenate_concurrent_files(outputDirNameConcurrent, outputFileName)
+        else:
+            if "outputFile" in locals():
+                outputFile.close()
+                print()
+                print('Done. Output file generated "%s".' % outputFileName)
+
 
 if __name__ == '__main__':
     main(sys.argv[1:])

tests/test_main.py

@@ -3,6 +3,7 @@
 
 import os
 import sys
+import time
 if sys.version_info[0] < 3:
     raise Exception("Python 2.x is not supported. Please upgrade to 3.x")
 
@@ -24,3 +25,18 @@ def test_QuerySearch():
     tweet = got.manager.TweetManager.getTweets(tweetCriteria)[0]
     assert '#europe' in tweet.hashtags.lower()
     assert '#refugees' in tweet.hashtags.lower()
+
+def test_MassFetchConcurrent():
+    time1 = time.time()
+    tweetCriteria = got.manager.TweetCriteria().setUsername("@realdonaldtrump").setMaxTweets(100).setSince("2018-01-01")
+    tweets1 = got.manager.ConcurrentTweetManager.getTweets(tweetCriteria,
+                                                           worker_count=5,
+                                                           forceMaxTweets=True,)
+    print("Time Needed Concurrent: {} Secs".format((time.time() - time1)))
+
+    time2 = time.time()
+    tweetCriteria = got.manager.TweetCriteria().setUsername("@realdonaldtrump").setMaxTweets(100).setSince("2018-01-01")
+    tweets2 = got.manager.TweetManager.getTweets(tweetCriteria)
+    print("Time Needed Non Concurrent: {} Secs".format((time.time() - time2)))
+
+    assert len(tweets1) == len(tweets2)