#!/usr/bin/env python """ Scrape Reddit posts and comments. CLI examples ------------ # Scrape data for a specific date python -m reddit_analysis.scraper.scrape --date 2025-04-20 """ from __future__ import annotations import argparse import os import sys from datetime import datetime, timedelta from pathlib import Path from typing import Optional, List, Dict, Any import pandas as pd import pyarrow.parquet as pq from huggingface_hub import ( hf_hub_download, list_repo_files, login, upload_file, HfApi ) import praw import logging import pytz from tqdm import tqdm from reddit_analysis.config_utils import setup_config class RedditAPI: """Wrapper class for Reddit API interactions that can be mocked for testing.""" def __init__(self, client_id: str, client_secret: str, user_agent: str): self.reddit = praw.Reddit( client_id=client_id, client_secret=client_secret, user_agent=user_agent ) def get_subreddit(self, name: str): return self.reddit.subreddit(name) def get_rate_limit_info(self) -> Dict[str, Any]: return { 'used': self.reddit.auth.limits.get('used'), 'remaining': self.reddit.auth.limits.get('remaining'), 'reset_timestamp': self.reddit.auth.limits.get('reset_timestamp') } class FileManager: """Wrapper class for file operations that can be mocked for testing.""" def __init__(self, base_dir: Path): self.base_dir = base_dir self.base_dir.mkdir(parents=True, exist_ok=True) def save_csv(self, df: pd.DataFrame, filename: str) -> Path: path = self.base_dir / f"{filename}.csv" df.to_csv(path, index=False) return path def save_parquet(self, df: pd.DataFrame, filename: str) -> Path: path = self.base_dir / f"{filename}.parquet" df.to_parquet(path, index=False) return path def read_parquet(self, filename: str) -> pd.DataFrame: path = self.base_dir / f"{filename}.parquet" return pd.read_parquet(path) class HuggingFaceManager: """Wrapper class for HuggingFace Hub operations that can be mocked for testing.""" def __init__(self, token: str, repo_id: str, repo_type: str = "dataset"): self.token = token self.repo_id = repo_id self.repo_type = repo_type self.api = HfApi(token=token) def download_file(self, path_in_repo: str) -> Path: return Path(hf_hub_download( repo_id=self.repo_id, repo_type=self.repo_type, filename=path_in_repo, token=self.token )) def upload_file(self, local_path: str, path_in_repo: str): self.api.upload_file( path_or_fileobj=local_path, path_in_repo=path_in_repo, repo_id=self.repo_id, repo_type=self.repo_type, token=self.token ) def list_files(self, prefix: str) -> List[str]: return self.api.list_repo_files( repo_id=self.repo_id, repo_type=self.repo_type ) class RedditScraper: def __init__( self, cfg: Dict[str, Any], reddit_api: Optional[RedditAPI] = None, file_manager: Optional[FileManager] = None, hf_manager: Optional[HuggingFaceManager] = None ): self.config = cfg['config'] self.secrets = cfg['secrets'] self.paths = cfg['paths'] self.logger = logging.getLogger(__name__) # Initialize services with dependency injection self.reddit_api = reddit_api or RedditAPI( client_id=self.secrets.get('REDDIT_CLIENT_ID'), client_secret=self.secrets.get('REDDIT_CLIENT_SECRET'), user_agent=self.secrets.get('REDDIT_USER_AGENT') ) self.file_manager = file_manager or FileManager(self.paths['raw_dir']) if self.config.get('push_to_hf', False): self.hf_manager = hf_manager or HuggingFaceManager( token=self.secrets.get('HF_TOKEN'), repo_id=self.config.get('repo_id'), repo_type=self.config.get('repo_type', 'dataset') ) else: self.hf_manager = hf_manager self.timezone = pytz.timezone(self.config['timezone']) self.logger.info(f"Output directory set to: {self.paths['raw_dir']}") def get_posts(self, subreddit_config: Dict[str, Any]) -> pd.DataFrame: """Fetch posts and comments from a subreddit based on configuration.""" subreddit_name = subreddit_config['name'] post_limit = subreddit_config['post_limit'] comment_limit = subreddit_config['comment_limit'] retrieved_at = datetime.now(self.timezone) records = [] subreddit = self.reddit_api.get_subreddit(subreddit_name) self.logger.info(f"Fetching {post_limit} posts from r/{subreddit_name}") for submission in tqdm( subreddit.top(time_filter="day", limit=post_limit), total=post_limit, desc=f"Processing r/{subreddit_name}" ): # Add post record records.append({ "subreddit": subreddit_name, "created_at": datetime.fromtimestamp(submission.created_utc, tz=self.timezone), "retrieved_at": retrieved_at, "type": "post", "text": submission.title + "\n\n" + submission.selftext, "score": submission.score, "post_id": submission.id, "parent_id": None }) # Get top comments if comment_limit > 0 if comment_limit > 0: submission.comment_sort = 'top' submission.comments.replace_more(limit=0) comments = getattr(submission.comments, '_comments', [])[:comment_limit] for comment in comments: records.append({ "subreddit": subreddit_name, "created_at": datetime.fromtimestamp(comment.created_utc, tz=self.timezone), "retrieved_at": retrieved_at, "type": "comment", "text": comment.body, "score": comment.score, "post_id": comment.id, "parent_id": comment.parent_id }) return pd.DataFrame(records) def print_rate_limit_info(self): """Print current Reddit API rate limit information.""" limits = self.reddit_api.get_rate_limit_info() reset_ts = limits.get('reset_timestamp') reset_time = ( datetime.fromtimestamp(reset_ts, tz=self.timezone) .strftime("%Y-%m-%d %I:%M:%S %p %Z") if reset_ts else "Unknown" ) self.logger.info("Reddit API Rate Limit Info") self.logger.info(f"Requests used: {limits.get('used')}") self.logger.info(f"Requests remaining: {limits.get('remaining')}") self.logger.info(f"Resets at: {reset_time}") def process_date(self, date_str: str) -> None: """Process data for a specific date.""" self.logger.info(f"Processing data for date: {date_str}") all_records = [] for sub_cfg in self.config['subreddits']: self.logger.info(f"Processing subreddit: {sub_cfg['name']}") df = self.get_posts(sub_cfg) all_records.append(df) combined_df = pd.concat(all_records, ignore_index=True) self.logger.info(f"Total records collected: {len(combined_df)}") # Save to CSV self.file_manager.save_csv(combined_df, date_str) # Upload to HuggingFace if configured if self.config.get('push_to_hf', False): self._upload_to_hf(combined_df, date_str) self.print_rate_limit_info() self.logger.info("Reddit scraper completed successfully") def _upload_to_hf(self, df: pd.DataFrame, date_str: str) -> None: """Upload data to HuggingFace Hub.""" try: current_date = datetime.strptime(date_str, "%Y-%m-%d") prev_date = (current_date - timedelta(days=1)).strftime("%Y-%m-%d") prev_file_path = f"{self.paths['hf_raw_dir']}/{prev_date}.parquet" self.logger.info(f"Checking for previous day's file: {prev_file_path}") try: downloaded_path = self.hf_manager.download_file(prev_file_path) existing_df = pd.read_parquet(downloaded_path) existing_ids = set(existing_df["post_id"].tolist()) Path(downloaded_path).unlink() original_count = len(df) df = df[~df["post_id"].isin(existing_ids)] filtered_count = len(df) self.logger.info(f"Filtered {original_count - filtered_count} duplicates") if df.empty: self.logger.info("No new posts to upload after deduplication") return except Exception as e: self.logger.warning(f"Could not fetch/process previous file: {e}") parquet_path = self.file_manager.save_parquet(df, date_str) path_in_repo = f"{self.paths['hf_raw_dir']}/{date_str}.parquet" self.hf_manager.upload_file(str(parquet_path), path_in_repo) self.logger.info(f"Uploaded {len(df)} rows for {date_str} → {path_in_repo}") except Exception as e: self.logger.error(f"Failed to upload to Hugging Face: {e}") raise def setup_logging(logs_dir: Path) -> logging.Logger: """Set up logging configuration using logs_dir from config.""" logs_dir.mkdir(parents=True, exist_ok=True) # Create log filename with current date log_file = logs_dir / f"reddit_scraper_{datetime.now().strftime('%Y-%m-%d')}.log" # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler(log_file, encoding="utf-8") ] ) logger = logging.getLogger(__name__) logger.info(f"Logging initialized. Log file: {log_file}") return logger def main(date_str: str = None) -> None: # Load configuration first cfg = setup_config() # Initialize logging with configured logs_dir logs_dir = cfg['paths']['logs_dir'] logger = setup_logging(logs_dir) logger.info("Starting Reddit scraper...") # Validate environment variables required_env_vars = ["REDDIT_CLIENT_ID", "REDDIT_CLIENT_SECRET", "REDDIT_USER_AGENT"] if cfg['config'].get('push_to_hf', False): required_env_vars.append("HF_TOKEN") missing = [v for v in required_env_vars if not cfg['secrets'].get(v) and not os.getenv(v)] if missing: logger.error(f"Missing required environment variables: {', '.join(missing)}") raise ValueError(f"Missing required environment variables: {', '.join(missing)}") # Instantiate and run logger.info("Initializing Reddit scraper...") scraper = RedditScraper(cfg) if date_str is None: date_str = datetime.now(pytz.timezone(cfg['config']['timezone'])).strftime("%Y-%m-%d") scraper.process_date(date_str) if __name__ == "__main__": from reddit_analysis.common_metrics import run_with_metrics parser = argparse.ArgumentParser(description='Scrape Reddit posts and comments.') parser.add_argument('--date', type=str, help='YYYY-MM-DD date to process') args = parser.parse_args() run_with_metrics("scrape", main, args.date)