datamule 1.6.0__tar.gz → 1.6.2__tar.gz

{datamule-1.6.0 → datamule-1.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datamule
-Version: 1.6.0
+Version: 1.6.2
 Summary: Work with SEC submissions at scale.
 Home-page: https://github.com/john-friedman/datamule-python
 Author: John Friedman

{datamule-1.6.0 → datamule-1.6.2}/datamule/portfolio.py

@@ -1,11 +1,13 @@
 from pathlib import Path
 from tqdm import tqdm
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from .submission import Submission
 from .sec.submissions.downloader import download as sec_download
 from .sec.submissions.textsearch import filter_text
 from .config import Config
 import os
+import tarfile
+from threading import Lock
 from .helper import _process_cik_and_metadata_filters
 from .seclibrary.downloader import download as seclibrary_download
 from .sec.xbrl.filter_xbrl import filter_xbrl
@@ -21,6 +23,10 @@ class Portfolio:
         self.submissions = []
         self.submissions_loaded = False
         self.MAX_WORKERS = os.cpu_count() - 1 
+        
+        # Batch tar support
+        self.batch_tar_handles = {}  # {batch_tar_path: tarfile_handle}
+        self.batch_tar_locks = {}    # {batch_tar_path: threading.Lock}
 
         self.monitor = Monitor()
         
@@ -34,9 +40,13 @@ class Portfolio:
         self.api_key = api_key
     
     def _load_submissions(self):
-        folders = [f for f in self.path.iterdir() if f.is_dir() or f.suffix=='.tar']
-        print(f"Loading {len(folders)} submissions")
+        print(f"Loading submissions")
+        
+        # Separate regular and batch items
+        regular_items = [f for f in self.path.iterdir() if (f.is_dir() or f.suffix=='.tar') and 'batch' not in f.name]
+        batch_tars = [f for f in self.path.iterdir() if f.is_file() and 'batch' in f.name and f.suffix == '.tar']
         
+        # Load regular submissions (existing logic)
         def load_submission(folder):
             try:
                 return Submission(folder)
@@ -44,17 +54,86 @@ class Portfolio:
                 print(f"Error loading submission from {folder}: {str(e)}")
                 return None
         
-        with ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
-            self.submissions = list(tqdm(
-                executor.map(load_submission, folders),
-                total=len(folders),
-                desc="Loading submissions"
-            ))
-            
-        # Filter out None values from failed submissions
-        self.submissions = [s for s in self.submissions if s is not None]
+        regular_submissions = []
+        if regular_items:
+            with ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
+                regular_submissions = list(tqdm(
+                    executor.map(load_submission, regular_items),
+                    total=len(regular_items),
+                    desc="Loading regular submissions"
+                ))
+        
+        # Load batch submissions with parallel processing + progress
+        batch_submissions = []
+        if batch_tars:
+            with tqdm(desc="Loading batch submissions", unit="submissions") as pbar:
+                with ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
+                    # Submit all batch tar jobs
+                    futures = [
+                        executor.submit(self._load_batch_submissions_worker, batch_tar, pbar) 
+                        for batch_tar in batch_tars
+                    ]
+                    
+                    # Collect results as they complete
+                    for future in as_completed(futures):
+                        try:
+                            batch_submissions.extend(future.result())
+                        except Exception as e:
+                            print(f"Error in batch processing: {str(e)}")
+        
+        # Combine and filter None values  
+        self.submissions = [s for s in (regular_submissions + batch_submissions) if s is not None]
         print(f"Successfully loaded {len(self.submissions)} submissions")
 
+    def _load_batch_submissions_worker(self, batch_tar_path, pbar):
+        """Worker function to load submissions from one batch tar with progress updates"""
+        try:
+            # Open tar handle and store it
+            tar_handle = tarfile.open(batch_tar_path, 'r')
+            self.batch_tar_handles[batch_tar_path] = tar_handle
+            self.batch_tar_locks[batch_tar_path] = Lock()
+            
+            # Find all accession directories
+            accession_prefixes = set()
+            for member in tar_handle.getmembers():
+                if '/' in member.name and member.name.endswith('metadata.json'):
+                    accession_prefix = member.name.split('/')[0]
+                    accession_prefixes.add(accession_prefix)
+            
+            # Create submissions for each accession
+            submissions = []
+            for accession_prefix in accession_prefixes:
+                try:
+                    submission = Submission(
+                        batch_tar_path=batch_tar_path,
+                        accession_prefix=accession_prefix,
+                        portfolio_ref=self
+                    )
+                    submissions.append(submission)
+                    pbar.update(1)  # Update progress for each successful submission
+                except Exception as e:
+                    print(f"Error loading batch submission {accession_prefix} from {batch_tar_path.name}: {str(e)}")
+            
+            return submissions
+            
+        except Exception as e:
+            print(f"Error loading batch tar {batch_tar_path}: {str(e)}")
+            return []
+
+    def _close_batch_handles(self):
+        """Close all open batch tar handles to free resources"""
+        for handle in self.batch_tar_handles.values():
+            try:
+                handle.close()
+            except Exception as e:
+                print(f"Error closing batch tar handle: {str(e)}")
+        self.batch_tar_handles.clear()
+        self.batch_tar_locks.clear()
+
+    def __del__(self):
+        """Cleanup batch tar handles on destruction"""
+        self._close_batch_handles()
+
     def process_submissions(self, callback):
         """Process all submissions using a thread pool."""
         if not self.submissions_loaded:
@@ -169,6 +248,7 @@ class Portfolio:
             )
 
         self.submissions_loaded = False
+        
     def monitor_submissions(self, data_callback=None, interval_callback=None,
                             polling_interval=1000, quiet=True, start_date=None,
                             validation_interval=600000):

{datamule-1.6.0 → datamule-1.6.2}/datamule/sec/submissions/monitor.py

@@ -9,16 +9,14 @@ from .eftsquery import EFTSQuery
 import aiohttp
 from zoneinfo import ZoneInfo 
 
-async def poll_rss(limiter):
+async def poll_rss(limiter, session):
     base_url = 'https://www.sec.gov/cgi-bin/browse-edgar?count=100&action=getcurrent&output=rss'
     
-    # Create a session specifically for this RSS polling operation
-    async with aiohttp.ClientSession(headers=headers) as session:
-        # Use the rate limiter before making the request
-        async with limiter:
-            # Make the HTTP request with the session
-            async with session.get(base_url) as response:
-                content = await response.read()
+    # Use the rate limiter before making the request
+    async with limiter:
+        # Use the provided session instead of creating a new one
+        async with session.get(base_url) as response:
+            content = await response.read()
     
     # Process the content
     content_str = content.decode('utf-8')
@@ -70,12 +68,31 @@ class Monitor():
         self.ratelimiters = {'sec.gov': PreciseRateLimiter(rate=5)}
         self.efts_query = EFTSQuery(quiet=True)
         self.efts_query.limiter = self.ratelimiters['sec.gov']
+        self.session = None
+        self.session_created_at = 0
+        self.session_lifetime = 300  # 5 minutes in seconds
 
     def set_domain_rate_limit(self, domain, rate):
         self.ratelimiters[domain] = PreciseRateLimiter(rate=rate)
         if domain == 'sec.gov':
             self.efts_query.limiter = self.ratelimiters[domain]
     
+    async def _ensure_fresh_session(self):
+        """Ensure we have a fresh session, recreating if expired or missing"""
+        current_time = time.time()
+        
+        # Check if we need a new session
+        if (self.session is None or 
+            current_time - self.session_created_at > self.session_lifetime):
+            
+            # Close old session if it exists
+            if self.session:
+                await self.session.close()
+            
+            # Create new session
+            self.session = aiohttp.ClientSession(headers=headers)
+            self.session_created_at = current_time
+    
     async def _async_run_efts_query(self, **kwargs):
         """Async helper method to run EFTS query without creating a new event loop"""
         # Make sure to set quiet parameter if provided in kwargs
@@ -103,83 +120,106 @@ class Monitor():
         if polling_interval is None and validation_interval is None:
             raise ValueError("At least one of polling_interval or validation_interval must be specified")
 
-        # Backfill if start_date is provided
-        if start_date is not None:
-            today_date = datetime.now(ZoneInfo("America/New_York")).strftime('%Y-%m-%d')
-            if not quiet:
-                print(f"Backfilling from {start_date} to {today_date}")
-
-            hits = clean_efts_hits(await self._async_run_efts_query(
-                filing_date=(start_date, today_date),
-                quiet=quiet
-            ))
-
-            new_hits = self._filter_new_accessions(hits)
-            if not quiet:
-                print(f"New submissions found: {len(new_hits)}")
-            if new_hits and data_callback:
-                data_callback(new_hits)
-
-        # Initialize timing variables
-        current_time = time.time()
-        last_polling_time = current_time
-        last_validation_time = current_time
-        
-        # Determine which operations to perform
-        do_polling = polling_interval is not None
-        do_validation = validation_interval is not None
+        # Ensure we have a fresh session
+        await self._ensure_fresh_session()
         
-        while True:
-            current_time = time.time()
-            
-            # RSS polling (if enabled)
-            if do_polling and (current_time - last_polling_time) >= polling_interval/1000:
-                if not quiet:
-                    print(f"Polling RSS feed")
-                results = await poll_rss(self.ratelimiters['sec.gov'])
-                new_results = self._filter_new_accessions(results)
-                if new_results:
-                    if not quiet:
-                        print(f"Found {len(new_results)} new submissions via RSS")
-                    if data_callback:
-                        data_callback(new_results)
-                last_polling_time = current_time
-            
-            # EFTS validation (if enabled)
-            if do_validation and (current_time - last_validation_time) >= validation_interval/1000:
-                # Get submissions from the last 24 hours for validation
+        try:
+            # Backfill if start_date is provided
+            if start_date is not None:
                 today_date = datetime.now(ZoneInfo("America/New_York")).strftime('%Y-%m-%d')
                 if not quiet:
-                    print(f"Validating submissions from {today_date}")
+                    print(f"Backfilling from {start_date} to {today_date}")
 
                 hits = clean_efts_hits(await self._async_run_efts_query(
-                    filing_date=(today_date, today_date),
+                    filing_date=(start_date, today_date),
                     quiet=quiet
                 ))
-                
+
                 new_hits = self._filter_new_accessions(hits)
-                if new_hits:
-                    if not quiet:
-                        print(f"Found {len(new_hits)} new submissions via EFTS validation")
-                    if data_callback:
-                        data_callback(new_hits)
-                last_validation_time = current_time
+                if not quiet:
+                    print(f"New submissions found: {len(new_hits)}")
+                if new_hits and data_callback:
+                    data_callback(new_hits)
+
+            # Initialize timing variables
+            current_time = time.time()
+            last_polling_time = current_time
+            last_validation_time = current_time
             
-            # Interval callback
-            if interval_callback:
-                interval_callback()
-
-            # Calculate next wake-up time
-            next_times = []
-            if do_polling:
-                next_times.append(last_polling_time + (polling_interval / 1000))
-            if do_validation:
-                next_times.append(last_validation_time + (validation_interval / 1000))
+            # Determine which operations to perform
+            do_polling = polling_interval is not None
+            do_validation = validation_interval is not None
             
-            next_wake_time = min(next_times)
-            current_time = time.time()
-            time_to_sleep = max(0, next_wake_time - current_time)
-            await asyncio.sleep(time_to_sleep)
+            while True:
+                current_time = time.time()
+                
+                # RSS polling (if enabled)
+                if do_polling and (current_time - last_polling_time) >= polling_interval/1000:
+                    if not quiet:
+                        print(f"Polling RSS feed")
+                    
+                    # Ensure session is fresh before polling
+                    await self._ensure_fresh_session()
+                    
+                    try:
+                        results = await poll_rss(self.ratelimiters['sec.gov'], self.session)
+                        new_results = self._filter_new_accessions(results)
+                        if new_results:
+                            if not quiet:
+                                print(f"Found {len(new_results)} new submissions via RSS")
+                            if data_callback:
+                                data_callback(new_results)
+                    except Exception as e:
+                        if not quiet:
+                            print(f"RSS polling error: {e}, will recreate session on next poll")
+                        # Force session recreation on next poll
+                        if self.session:
+                            await self.session.close()
+                            self.session = None
+                    
+                    last_polling_time = current_time
+                
+                # EFTS validation (if enabled)
+                if do_validation and (current_time - last_validation_time) >= validation_interval/1000:
+                    # Get submissions from the last 24 hours for validation
+                    today_date = datetime.now(ZoneInfo("America/New_York")).strftime('%Y-%m-%d')
+                    if not quiet:
+                        print(f"Validating submissions from {today_date}")
+
+                    hits = clean_efts_hits(await self._async_run_efts_query(
+                        filing_date=(today_date, today_date),
+                        quiet=quiet
+                    ))
+                    
+                    new_hits = self._filter_new_accessions(hits)
+                    if new_hits:
+                        if not quiet:
+                            print(f"Found {len(new_hits)} new submissions via EFTS validation")
+                        if data_callback:
+                            data_callback(new_hits)
+                    last_validation_time = current_time
+                
+                # Interval callback
+                if interval_callback:
+                    interval_callback()
+
+                # Calculate next wake-up time
+                next_times = []
+                if do_polling:
+                    next_times.append(last_polling_time + (polling_interval / 1000))
+                if do_validation:
+                    next_times.append(last_validation_time + (validation_interval / 1000))
+                
+                next_wake_time = min(next_times)
+                current_time = time.time()
+                time_to_sleep = max(0, next_wake_time - current_time)
+                await asyncio.sleep(time_to_sleep)
+                
+        finally:
+            # Clean up the session when done
+            if self.session:
+                await self.session.close()
+                self.session = None
 
     def monitor_submissions(self, data_callback=None, interval_callback=None,
                             polling_interval=1000, quiet=True, start_date=None,