crawl_follows.py scripts/crawl_follows.py

+1

data/likes/.gitignore

··· 1 1 + ./*

get_posts.py scripts/get_posts.py

+259

scripts/get_likes.py

··· 1 1 + import argparse 2 2 + import asyncio 3 3 + from datetime import datetime 4 4 + import gzip 5 5 + import json 6 6 + import logging 7 7 + import os 8 8 + import sys 9 9 + from typing import Tuple, List, Dict 10 10 + 11 11 + from atproto import AsyncClient 12 12 + from atproto import exceptions as at_exceptions 13 13 + from atproto_client.models.app.bsky.feed.like import Record as LikeRecord 14 14 + import pandas as pd 15 15 + 16 16 + from crawl_follows import RateLimit 17 17 + 18 18 + logger = logging.getLogger(__name__) 19 19 + logger.setLevel(logging.INFO) 20 20 + 21 21 + # Create formatter 22 22 + formatter = logging.Formatter( 23 23 + "%(asctime)s | %(levelname)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S" 24 24 + ) 25 25 + 26 26 + # Console handler 27 27 + console_handler = logging.StreamHandler(sys.stdout) 28 28 + console_handler.setFormatter(formatter) 29 29 + logger.addHandler(console_handler) 30 30 + 31 31 + 32 32 + BATCH_SIZE = 1 33 33 + FOLLOWER_THRESHOLD = 150 34 34 + REQUIRED_ENV = ("BSKY_USER", "BSKY_APP_PW") 35 35 + 36 36 + 37 37 + async def get_all_likes( 38 38 + client: AsyncClient, 39 39 + rate_limit: RateLimit, 40 40 + account_did: str, 41 41 + start_dt: datetime, 42 42 + end_dt: datetime, 43 43 + ) -> Tuple[List[LikeRecord], str]: 44 44 + likes: List[LikeRecord] = [] 45 45 + await rate_limit.acquire() 46 46 + try: 47 47 + data = await client.com.atproto.repo.list_records( 48 48 + { 49 49 + "collection": "app.bsky.feed.like", 50 50 + "repo": account_did, 51 51 + "limit": 100, 52 52 + } 53 53 + ) 54 54 + # If user can't be accessed just return an empty list to skip next time 55 55 + except at_exceptions.BadRequestError as e: 56 56 + if e.response.status_code == 400: 57 57 + return [], account_did 58 58 + else: 59 59 + logger.info(f"Error status code: {e.response.status_code}") 60 60 + raise e 61 61 + 62 62 + for like in data.records: 63 63 + dt = datetime.strptime(like.value.created_at, "%Y-%m-%dT%H:%M:%S.%fZ") 64 64 + if start_dt <= dt and dt < end_dt: 65 65 + likes.append(like) 66 66 + 67 67 + hit_start_window = False 68 68 + while data.cursor and not hit_start_window: 69 69 + await rate_limit.acquire() 70 70 + data = await client.com.atproto.repo.list_records( 71 71 + { 72 72 + "collection": "app.bsky.feed.like", 73 73 + "repo": account_did, 74 74 + "cursor": data.cursor, 75 75 + "limit": 100, 76 76 + } 77 77 + ) 78 78 + 79 79 + for like in data.records: 80 80 + dt = datetime.strptime(like.value.created_at, "%Y-%m-%dT%H:%M:%S.%fZ") 81 81 + if start_dt <= dt and dt < end_dt: 82 82 + likes.append(like) 83 83 + 84 84 + if dt < start_dt: 85 85 + hit_start_window = True 86 86 + 87 87 + return likes, account_did 88 88 + 89 89 + 90 90 + async def retrieve_likes( 91 91 + user: str, 92 92 + app_pw: str, 93 93 + graph_file: str, 94 94 + checkpoint_dir: str, 95 95 + start_dt: datetime, 96 96 + end_dt: datetime, 97 97 + ): 98 98 + 99 99 + # If checkpoint dir doesn't exist, try to create it 100 100 + if not os.path.isdir(checkpoint_dir): 101 101 + logger.info("Checkpoint dir doesn't exist, creating...") 102 102 + try: 103 103 + os.mkdir(checkpoint_dir) 104 104 + except Exception as e: 105 105 + logger.error(f"Failed to created checkpoint dir, {checkpoint_dir}\n{e}") 106 106 + sys.exit(1) 107 107 + 108 108 + # Checkpoint folders contain one file per user 109 109 + completed_accounts = set() 110 110 + try: 111 111 + files = os.listdir(checkpoint_dir) 112 112 + for file in files: 113 113 + # Grab entire file name except for .gz extension 114 114 + completed_accounts.add(file[:-3]) 115 115 + except Exception as e: 116 116 + logger.error( 117 117 + f"Failed to recover from checkpoint dir, {checkpoint_dir}\n{e}", 118 118 + exc_info=1, 119 119 + ) 120 120 + sys.exit(1) 121 121 + 122 122 + # Load follow graph parquet file 123 123 + to_explore = dict() 124 124 + try: 125 125 + logger.info("Parsing follower graph file...") 126 126 + follow_df = pd.read_parquet(graph_file) 127 127 + # Limit to only accounts following between 100 and 1000 followers 128 128 + follow_df = follow_df.loc[follow_df["follows"].str.len().between(100, 1000)] 129 129 + except Exception as e: 130 130 + logger.error(f"Failed to open follow graph file, {graph_file}\n{e}") 131 131 + sys.exit(1) 132 132 + 133 133 + for _, row in follow_df.iterrows(): 134 134 + for acct in row["follows"]: 135 135 + if acct not in completed_accounts: 136 136 + if acct not in to_explore: 137 137 + to_explore[acct] = 0 138 138 + to_explore[acct] += 1 139 139 + 140 140 + accts = [(acct, follows) for acct, follows in to_explore.items()] 141 141 + accts.sort(key=lambda x: -1 * x[1]) 142 142 + 143 143 + logger.info(f"Num of accounts to retrieve posts from: {len(accts)}") 144 144 + 145 145 + client = AsyncClient() 146 146 + await client.login(user, app_pw) 147 147 + 148 148 + # Get all posts for accounts 149 149 + batch_count = 0 150 150 + fail_count = 0 151 151 + rate_limiter = RateLimit(BATCH_SIZE) 152 152 + for i in range(0, len(accts), BATCH_SIZE): 153 153 + batch = [acct for acct, _ in accts[i : i + BATCH_SIZE]] 154 154 + for result in asyncio.as_completed( 155 155 + [ 156 156 + get_all_likes(client, rate_limiter, did, start_dt, end_dt) 157 157 + for did in batch 158 158 + ] 159 159 + ): 160 160 + try: 161 161 + likes, did = await result 162 162 + # Save likes 163 163 + with gzip.open( 164 164 + os.path.join(checkpoint_dir, did + ".gz"), "wt" 165 165 + ) as out_file: 166 166 + for like in likes: 167 167 + out_file.write(like.model_dump_json() + "\n") 168 168 + except at_exceptions.BadRequestError as e: 169 169 + # Bad request is probably a profile that's private or deleted 170 170 + logger.info(f"Bad Request: {e.response.content.error}") 171 171 + continue 172 172 + except Exception as e: 173 173 + logger.error(f"Failed to get likes: {e}", exc_info=1) 174 174 + fail_count += 1 175 175 + if fail_count >= 100: 176 176 + logger.error("Hitting error threshold, exiting...") 177 177 + sys.exit(1) 178 178 + continue 179 179 + 180 180 + batch_count += 1 181 181 + if batch_count % 10 == 0: 182 182 + logger.info(f"Completed batch: {batch_count}") 183 183 + sys.exit(0) 184 184 + 185 185 + 186 186 + def main(): 187 187 + for key in REQUIRED_ENV: 188 188 + if key not in os.environ: 189 189 + raise ValueError(f"Must set '{key}' env var") 190 190 + 191 191 + user_name = os.environ["BSKY_USER"] 192 192 + app_pw = os.environ["BSKY_APP_PW"] 193 193 + 194 194 + parser = argparse.ArgumentParser( 195 195 + prog="GetLikes", 196 196 + description="Get all likes for accounts in provided follow graph", 197 197 + ) 198 198 + parser.add_argument( 199 199 + "--graph-file", 200 200 + dest="graph_file", 201 201 + required=True, 202 202 + help="File with follow graph", 203 203 + ) 204 204 + parser.add_argument( 205 205 + "--save-dir", 206 206 + dest="save_dir", 207 207 + required=True, 208 208 + help="Where to store crawl data", 209 209 + ) 210 210 + parser.add_argument( 211 211 + "--start", 212 212 + dest="start", 213 213 + required=True, 214 214 + help="Date to start saving posts from (YYYY-MM-DD)", 215 215 + ) 216 216 + parser.add_argument( 217 217 + "--end", 218 218 + dest="end", 219 219 + required=True, 220 220 + help="Date to stop (exclusive) saving posts from (YYYY-MM-DD)", 221 221 + ) 222 222 + args = parser.parse_args() 223 223 + 224 224 + if args.save_dir is None and args.ckpt is None: 225 225 + logger.error("Must provide save dir or checkpoint dir") 226 226 + sys.exit(1) 227 227 + 228 228 + try: 229 229 + start = datetime.strptime(args.start, "%Y-%m-%d") 230 230 + except: 231 231 + logger.error("Invalid start date") 232 232 + sys.exit(1) 233 233 + 234 234 + try: 235 235 + end = datetime.strptime(args.end, "%Y-%m-%d") 236 236 + except: 237 237 + logger.error("Invalid end date") 238 238 + sys.exit(1) 239 239 + 240 240 + if end <= start: 241 241 + logger.error( 242 242 + "Start date has to be before date, what're you trying to do man..." 243 243 + ) 244 244 + sys.exit(1) 245 245 + 246 246 + asyncio.run( 247 247 + retrieve_likes( 248 248 + user_name, 249 249 + app_pw, 250 250 + graph_file=args.graph_file, 251 251 + checkpoint_dir=args.save_dir, 252 252 + start_dt=start, 253 253 + end_dt=end, 254 254 + ) 255 255 + ) 256 256 + 257 257 + 258 258 + if __name__ == "__main__": 259 259 + main()