diff --git a/src/process/twitter_process.py b/src/process/twitter_process.py index 93c6371..bee7017 100644 --- a/src/process/twitter_process.py +++ b/src/process/twitter_process.py @@ -1,7 +1,10 @@ +import json import os from dataclasses import dataclass from datetime import datetime +from utils import normalize_directory + @dataclass class GeneralUser: @@ -29,3 +32,31 @@ class Posting: date: datetime +def load_users_popularity(user_dir: str = './data/twitter/user/') -> list[tuple[str, int]]: + """ + After downloading a wide range of users using download_users_start in raw_collect/twitter.py, + this function will read the user files and rank the users by popularity. + + The return format will consist of a list of users' screen names and popularity. + + :param user_dir: Download directory of users data, should be the same as the downloads dir in + download_user_start. (Default: "./data/twitter/user/") + :return: List of users' screen names and popularity, sorted descending by popularity. + """ + user_dir = normalize_directory(user_dir) + users = [] + + # Loop through all the files + for filename in os.listdir(user_dir + ''): + # Only check json files and ignore macos dot files + if filename.endswith('.json') and not filename.startswith('.'): + # Read + with open(filename, 'r', encoding='utf-8') as f: + user = json.load(f) + users.append((user['screen_name'], user['followers_count'])) + + # Sort by followers count, descending + users.sort(key=lambda x: x[1]) + + # Return + return users