# TODO: move queue structures and file handling here import os import shutil import traceback import util import talenttweet as tt # User timestamps line format: # {user_id} {status_num} {UNIX_timestamp} class TalentTweetQueue: instance = None def __init__(self): TalentTweetQueue.instance = self self.queue_path = util.get_queue_path() self.queue_backup_path = util.get_queue_backup_path() self.current_ttweet_path = f'{util.get_project_dir()}/_current_ttweet.txt' self.finished_ttweets_path = f'{util.get_project_dir()}/finished_ttweets.txt' self.is_good = True self.__sorted = False self.finished_user_dates = dict() self.ttweets_dict = dict() self.finished_ttweets = list() ## file check, backup copy if os.path.exists(self.queue_backup_path): print('Found backup queue! We errored in the previous run.') shutil.copyfile(self.queue_backup_path, self.queue_path) elif os.path.exists(self.queue_path): print('Creating backup queue...') shutil.copyfile(self.queue_path, self.queue_backup_path) ## initialize structures # user timestamps try: with open(self.queue_path, 'r') as f: for line in f: tokens = line.split() if len(tokens) == 0: continue if tokens[0][0] != '#': print(f'Stopped finding user dates at {line}') # reached end of accounts list break if tokens[2] != '-1': self.finished_user_dates[int(tokens[1])] = tokens[2] except: pass # ttweets try: with open(self.queue_path, 'r') as f: # reset seek head # Get existing queued TalentTweets for line in f: tokens = line.split() if len(tokens) == 0 or tokens[0][0] == '#': continue ttweet = tt.TalentTweet.deserialize(line) # print(f'{ttweet.tweet_id}:\n{ttweet}') self.ttweets_dict[ttweet.tweet_id] = ttweet print(f'Found {len(self.finished_user_dates)} scraped accounts and {len(self.ttweets_dict)} tweets in queue.') except: traceback.print_exc() pass # finished ttweets try: with open(self.finished_ttweets_path, 'r') as f: for line in f: self.finished_ttweets.append(int(line)) except: pass def is_empty(self): return self.get_count() <= 0 def add_ttweet(self, ttweet): self.ttweets_dict[ttweet.tweet_id] = ttweet self.__sorted = False def get_ttweet(self, id): return self.ttweets_dict[id] def get_next_ttweet(self): self.is_good = False if os.path.exists(self.current_ttweet_path): with open(self.current_ttweet_path, 'r') as f: ttweet = tt.TalentTweet.deserialize(f.readline()) if ttweet.tweet_id in self.ttweets_dict: self.ttweets_dict.pop(ttweet.tweet_id) return ttweet self.__sort_ttweets_dict() key = list(self.ttweets_dict.keys())[0] ttweet = self.ttweets_dict.pop(key) with open(self.current_ttweet_path, 'w') as f: f.write(ttweet.serialize()) return ttweet def get_count(self): return len(self.ttweets_dict) ## Call when the TalentTweet retrieved from get_next_ttweet() was # posted successfully. def good(self): with open(self.current_ttweet_path, 'r') as f: ttweet = tt.TalentTweet.deserialize(f.readline()) self.add_finished_tweet(ttweet.tweet_id) os.remove(self.current_ttweet_path) self.save_file() self.is_good = True # overwrite queue.txt def save_file(self): print('saving file...', end='') shutil.copyfile(self.queue_path, self.queue_backup_path) self.__sort_ttweets_dict() with open(self.queue_path, 'w') as f: # write dates for (id, date) in self.finished_user_dates.items(): f.write(f'# {id} {date}\n') f.write('\n') # write sorted ttweets for ttweet in self.ttweets_dict.values(): f.write(ttweet.serialize() + '\n') print('done') def add_finished_tweet(self, id): self.finished_ttweets.append(id) with open(self.finished_ttweets_path, 'a') as f: f.write(f'{id}\n') def __sort_ttweets_dict(self): if not self.__sorted: self.ttweets_dict = dict(sorted(self.ttweets_dict.items())) self.__sorted = True # destructor def __del__(self): if self.is_good: print('Ended in good state, deleting backup queue...') os.remove(self.queue_backup_path) else: print('Ended in bad state, keeping backup queue.')