Resolved Discussions

Author: giulionf

GetOldTweets3/manager/ConcurrentTweetManager.py

@@ -1,8 +1,8 @@
 # Author: Giulio Neusch-Frediani - www.github.com/giulionf
 from datetime import datetime
+from datetime import timedelta
 from threading import Thread
-from queue import Queue
-from queue import Empty
+from threading import Lock
 import copy
 
 from GetOldTweets3.manager.TweetManager import TweetManager
@@ -14,42 +14,61 @@
 class ConcurrentTweetManager:
 
     @staticmethod
-    def getTweets(tweetCriteria, receiveBuffer=None, bufferLength=100, proxy=None, debug=False, worker_count=1, ordered=True):
+    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 = Queue()
+        time_spans = []
         tweets = []
         workers = []
+        lock = Lock()
 
         # Split the date in smaller parts
-        since_date = datetime.strptime(tweetCriteria.since, DATE_FORMAT) if tweetCriteria.since else FIRST_TWEET_DATE
-        until_date = datetime.strptime(tweetCriteria.until, DATE_FORMAT) if tweetCriteria.until else datetime.now()
+        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)
-        criteria.setMaxTweets(criteria.maxTweets / worker_count)
+        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.put((from_time, to_time))
+            time_spans.append((from_time, to_time))
 
-            w = WorkerThread(copy.deepcopy(criteria), time_spans, tweets, receiveBuffer, bufferLength, proxy, debug)
+        for i in range(0, worker_count):
+            w = WorkerThread(copy.deepcopy(criteria), time_spans[i], tweets, lock, receiveBuffer, bufferLength,
+                             proxy, debug)
             w.start()
             workers.append(w)
 
-        # Wait for the workers to finish and then stop them
-        time_spans.join()
+        # Wait for the workers to finish, then return the results
         for worker in workers:
-            worker.stop()
+            worker.join()
 
         return sorted(tweets, key=lambda r: r.date)
 
 
 class WorkerThread(Thread):
 
-    def __init__(self, tweetCriteria, time_spans, tweets, receiveBuffer=None, bufferLength=100, proxy=None, debug=False):
+    def __init__(self, tweetCriteria, time_span, tweets, lock,  receiveBuffer=None, bufferLength=100, proxy=None,
+                 debug=False):
         super().__init__()
         self.stopped = False
         self.manager = TweetManager()
@@ -58,21 +77,13 @@ def __init__(self, tweetCriteria, time_spans, tweets, receiveBuffer=None, buffer
         self.bufferLength = bufferLength
         self.proxy = proxy
         self.debug = debug
-        self.time_spans = time_spans
+        self.time_span = time_span
         self.tweets = tweets
-
-    def stop(self):
-        self.stopped = True
+        self.lock = lock
 
     def run(self) -> None:
-        while not self.stopped:
-            try:
-                time_span = self.time_spans.get_nowait()
-                self.tweetCriteria.setSince(datetime.strftime(time_span[0], "%Y-%m-%d"))
-                self.tweetCriteria.setUntil(datetime.strftime(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)
-                self.time_spans.task_done()
-            except Empty:
-                pass
+        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)

bin/GetOldTweets3

@@ -43,6 +43,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.')
@@ -64,12 +82,18 @@ def main(argv):
                                               "maxtweets=",
                                               "lang=",
                                               "output=",
-                                              "debug"))
+                                              "debug",
+                                              "workers=",
+                                              "forcemaxtweets"))
 
         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:
@@ -107,10 +131,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__)
@@ -136,8 +167,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):
@@ -167,8 +202,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')
@@ -182,10 +254,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

@@ -27,19 +27,18 @@ def test_QuerySearch():
 
 def test_MassFetchConcurrent():
     time1 = time.time()
-    tweetCriteria = got.manager.TweetCriteria().setQuerySearch('#europe #refugees')\
-                                               .setSince("2015-05-01")\
-                                               .setUntil("2015-09-30")\
-                                               .setMaxTweets(200)
-    tweets1 = got.manager.ConcurrentTweetManager.getTweets(tweetCriteria, worker_count=5)
+    tweetCriteria = got.manager.TweetCriteria().setUsername("@realdonaldtrump")
+    tweets1 = got.manager.ConcurrentTweetManager.getTweets(tweetCriteria, worker_count=100, forceMaxTweets=True,
+                                                           receiveBuffer=on_receive_batch)
     print("Time Needed Concurrent: {} Secs".format((time.time() - time1)))
 
     time2 = time.time()
-    tweetCriteria = got.manager.TweetCriteria().setQuerySearch('#europe #refugees')\
-                                               .setSince("2015-05-01")\
-                                               .setUntil("2015-09-30")\
-                                               .setMaxTweets(200)
-    tweets2 = got.manager.TweetManager.getTweets(tweetCriteria)
+    tweetCriteria = got.manager.TweetCriteria().setUsername("@realdonaldtrump")
+    tweets2 = got.manager.TweetManager.getTweets(tweetCriteria, receiveBuffer=on_receive_batch)
     print("Time Needed Non Concurrent: {} Secs".format((time.time() - time2)))
 
     assert len(tweets1) == len(tweets2)
+
+
+def on_receive_batch(results_aux):
+    print(str(len(results_aux)))