Initial database setup.

main.py

@@ -3,19 +3,22 @@
 import argparse
 import json
 import os
-from collections import Counter
+import time
 
 import praw
 import yfinance as yf
 from dotenv import load_dotenv
 
+import database
 from ticker_extractor import extract_tickers
 
 # Load environment variables from .env file
 load_dotenv()
+# How old (in seconds) market cap data can be before we refresh it. 24 hours = 86400 seconds.
+MARKET_CAP_REFRESH_INTERVAL = 86400
 
 def load_subreddits(filepath):
-    """Loads a list of subreddits from a JSON file."""
+    # (This function is unchanged)
     try:
         with open(filepath, 'r') as f:
             data = json.load(f)
@@ -28,21 +31,17 @@ def load_subreddits(filepath):
         return None
 
 def get_market_cap(ticker_symbol):
-    """Fetches the market capitalization for a given stock ticker."""
+    """Fetches the market capitalization for a given stock ticker from yfinance."""
     try:
         ticker = yf.Ticker(ticker_symbol)
-        market_cap = ticker.info.get('marketCap')
-        if market_cap:
-            # Formatting for better readability
-            return f"${market_cap:,}"
-
-        return "N/A"
-    except Exception as e:
-        # yfinance can sometimes fail for various reasons (e.g., invalid ticker)
-        return "N/A"
+        # .info can be slow; .fast_info is a lighter alternative
+        market_cap = ticker.fast_info.get('marketCap')
+        return market_cap if market_cap else None
+    except Exception:
+        return None
 
 def get_reddit_instance():
-    """Initializes and returns a PRAW Reddit instance."""
+    # (This function is unchanged)
     client_id = os.getenv("REDDIT_CLIENT_ID")
     client_secret = os.getenv("REDDIT_CLIENT_SECRET")
     user_agent = os.getenv("REDDIT_USER_AGENT")
@@ -50,41 +49,54 @@ def get_reddit_instance():
     if not all([client_id, client_secret, user_agent]):
         print("Error: Reddit API credentials not found in .env file.")
         return None
+    return praw.Reddit(client_id=client_id, client_secret=client_secret, user_agent=user_agent)
 
-    return praw.Reddit(
-        client_id=client_id,
-        client_secret=client_secret,
-        user_agent=user_agent
-    )
 
 def scan_subreddits(reddit, subreddits_list, post_limit=25):
-    """Scans subreddits for stock tickers and returns a count of each."""
-    all_tickers = Counter()
-
+    """Scans subreddits, stores mentions, and updates market caps in the database."""
+    conn = database.get_db_connection()
+    
     print(f"\nScanning {len(subreddits_list)} subreddits for top {post_limit} posts...")
     for subreddit_name in subreddits_list:
         try:
+            subreddit_id = database.get_or_create_entity(conn, 'subreddits', 'name', subreddit_name)
             subreddit = reddit.subreddit(subreddit_name)
-            print(f"r/{subreddit_name}...")
-            # Fetch hot posts from the subreddit
+            print(f"Scanning r/{subreddit_name}...")
+            
             for submission in subreddit.hot(limit=post_limit):
-                # Combine title and selftext for analysis
                 full_text = submission.title + " " + submission.selftext
-
-                # Extract tickers from the combined text
                 tickers_in_post = extract_tickers(full_text)
-                all_tickers.update(tickers_in_post)
-
-                # Future work: also scan comments
-                # submission.comments.replace_more(limit=0) # Expand all comment trees
-                # for comment in submission.comments.list():
-                #     tickers_in_comment = extract_tickers(comment.body)
-                #     all_tickers.update(tickers_in_comment)
+                
+                for ticker_symbol in set(tickers_in_post):
+                    ticker_id = database.get_or_create_entity(conn, 'tickers', 'symbol', ticker_symbol)
+                    
+                    database.add_mention(
+                        conn,
+                        ticker_id=ticker_id,
+                        subreddit_id=subreddit_id,
+                        post_id=submission.id,
+                        timestamp=int(submission.created_utc)
+                    )
+                    
+                    # --- Check if market cap needs updating ---
+                    ticker_info = database.get_ticker_info(conn, ticker_id)
+                    current_time = int(time.time())
+                    
+                    if not ticker_info['last_updated'] or (current_time - ticker_info['last_updated'] > MARKET_CAP_REFRESH_INTERVAL):
+                        print(f"  -> Fetching market cap for {ticker_symbol}...")
+                        market_cap = get_market_cap(ticker_symbol)
+                        if market_cap:
+                            database.update_ticker_market_cap(conn, ticker_id, market_cap)
+                        else:
+                            # If fetch fails, still update the timestamp so we don't try again for 24 hours
+                            database.update_ticker_market_cap(conn, ticker_id, ticker_info['market_cap']) # Keep old value
 
         except Exception as e:
             print(f"Could not scan r/{subreddit_name}. Error: {e}")
+            
+    conn.close()
+    print("\n--- Scan Complete ---")
 
-    return all_tickers
 
 def main():
     """Main function to run the Reddit stock analysis tool."""
@@ -92,28 +104,21 @@ def main():
     parser.add_argument("config_file", help="Path to the JSON file containing subreddits.")
     args = parser.parse_args()
 
-    # --- Part 1: Load Configuration & Initialize Reddit ---
+    # --- Part 1: Initialize ---
+    database.initialize_db()
+    
     subreddits = load_subreddits(args.config_file)
-    if not subreddits:
-        return
+    if not subreddits: return
 
     reddit = get_reddit_instance()
-    if not reddit:
-        return
+    if not reddit: return
 
-    # --- Part 2: Scan Reddit for Tickers ---
-    ticker_counts = scan_subreddits(reddit, subreddits)
-    if not ticker_counts:
-        print("No tickers found.")
-        return
+    # --- Part 2: Scan and Store ---
+    scan_subreddits(reddit, subreddits)
+    
+    # --- Part 3: Generate and Display Report ---
+    database.generate_summary_report()
 
-    print("\n--- Scan Complete ---")
-    print("Top 15 mentioned tickers:")
-
-    # --- Part 3: Display Results ---
-    # We will enrich this data with market cap and sentiment in the next steps
-    for ticker, count in ticker_counts.most_common(15):
-        print(f"{ticker}: {count} mentions")
 
 if __name__ == "__main__":
     main()