datamule 0.422__cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

datamule/downloader/downloader.py

@@ -0,0 +1,364 @@
+import asyncio
+import aiohttp
+import os
+from tqdm import tqdm
+from datetime import datetime
+from urllib.parse import urlencode
+import aiofiles
+import json
+import time
+from collections import deque
+
+from ..helper import identifier_to_cik, load_package_csv, fix_filing_url, headers
+from ..parser.sgml_parsing.sgml_parser_cy import parse_sgml_submission
+
+class RetryException(Exception):
+    def __init__(self, url, retry_after=601):
+        self.url = url
+        self.retry_after = retry_after
+
+class PreciseRateLimiter:
+    def __init__(self, rate, interval=1.0):
+        self.rate = rate  # requests per interval
+        self.interval = interval  # in seconds
+        self.token_time = self.interval / self.rate  # time per token
+        self.last_time = time.time()
+        self.lock = asyncio.Lock()
+
+    async def acquire(self):
+        async with self.lock:
+            now = time.time()
+            wait_time = self.last_time + self.token_time - now
+            if wait_time > 0:
+                await asyncio.sleep(wait_time)
+            self.last_time = time.time()
+            return True
+
+    async def __aenter__(self):
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+class RateMonitor:
+    def __init__(self, window_size=1.0):
+        self.window_size = window_size
+        self.requests = deque()
+        self._lock = asyncio.Lock()
+    
+    async def add_request(self, size_bytes):
+        async with self._lock:
+            now = time.time()
+            self.requests.append((now, size_bytes))
+            while self.requests and self.requests[0][0] < now - self.window_size:
+                self.requests.popleft()
+    
+    def get_current_rates(self):
+        now = time.time()
+        while self.requests and self.requests[0][0] < now - self.window_size:
+            self.requests.popleft()
+        
+        if not self.requests:
+            return 0, 0
+        
+        request_count = len(self.requests)
+        byte_count = sum(size for _, size in self.requests)
+        
+        requests_per_second = request_count / self.window_size
+        mb_per_second = (byte_count / 1024 / 1024) / self.window_size
+        
+        return round(requests_per_second, 1), round(mb_per_second, 2)
+
+class Downloader:
+    def __init__(self):
+        self.headers = headers
+        self.limiter = PreciseRateLimiter(5)  # 10 requests per second
+        self.session = None
+        self.parse_filings = True
+        self.download_queue = asyncio.Queue()
+        self.rate_monitor = RateMonitor()
+        self.current_pbar = None
+        self.connection_semaphore = asyncio.Semaphore(5)
+    
+    def update_progress_description(self):
+        if self.current_pbar:
+            reqs_per_sec, mb_per_sec = self.rate_monitor.get_current_rates()
+            self.current_pbar.set_description(
+                f"Progress [Rate: {reqs_per_sec}/s | {mb_per_sec} MB/s]"
+            )
+
+    async def __aenter__(self):
+        await self._init_session()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._close()
+
+    async def _init_session(self):
+        if not self.session:
+            self.session = aiohttp.ClientSession(headers=self.headers)
+
+    async def _close(self):
+        if self.session:
+            await self.session.close()
+            self.session = None
+
+    async def _fetch_json(self, url):
+        """Fetch JSON with rate monitoring."""
+        async with self.limiter:
+            try:
+                url = fix_filing_url(url)
+                async with self.session.get(url) as response:
+                    if response.status == 429:
+                        raise RetryException(url)
+                    response.raise_for_status()
+                    content = await response.read()
+                    await self.rate_monitor.add_request(len(content))
+                    self.update_progress_description()
+                    return await response.json()
+            except aiohttp.ClientResponseError as e:
+                if e.status == 429:
+                    raise RetryException(url)
+                raise
+
+    async def _get_filing_urls_from_efts(self, base_url):
+        """Fetch filing URLs from EFTS in batches."""
+        start = 0
+        page_size = 100
+        urls = []
+        
+        data = await self._fetch_json(f"{base_url}&from=0&size=1")
+        if not data or 'hits' not in data:
+            return []
+            
+        total_hits = data['hits']['total']['value']
+        if not total_hits:
+            return []
+
+        pbar = tqdm(total=total_hits, desc="Fetching URLs [Rate: 0/s | 0 MB/s]")
+        self.current_pbar = pbar
+        
+        while start < total_hits:
+            try:
+                tasks = [
+                    self._fetch_json(f"{base_url}&from={start + i * page_size}&size={page_size}") 
+                    for i in range(10)
+                ]
+                
+                results = await asyncio.gather(*tasks)
+                
+                for data in results:
+                    if data and 'hits' in data:
+                        hits = data['hits']['hits']
+                        if hits:
+                            batch_urls = [
+                                f"https://www.sec.gov/Archives/edgar/data/{hit['_source']['ciks'][0]}/{hit['_id'].split(':')[0]}.txt" 
+                                for hit in hits
+                            ]
+                            urls.extend(batch_urls)
+                            pbar.update(len(hits))
+                            self.update_progress_description()
+                
+                start += 10 * page_size
+
+            except RetryException as e:
+                print(f"\nRate limited. Sleeping for {e.retry_after} seconds...")
+                await asyncio.sleep(e.retry_after)
+                continue
+            except Exception as e:
+                print(f"\nError fetching URLs batch at {start}: {str(e)}")
+                break
+
+        pbar.close()
+        self.current_pbar = None
+        return urls
+
+    async def _download_file(self, url, filepath):
+        """Download single file with precise rate limiting."""
+        async with self.connection_semaphore:
+            async with self.limiter:
+                try:
+                    url = fix_filing_url(url)
+                    async with self.session.get(url) as response:
+                        if response.status == 429:
+                            raise RetryException(url)
+                        response.raise_for_status()
+                        content = await response.read()
+                        await self.rate_monitor.add_request(len(content))
+                        self.update_progress_description()
+                        
+                        parsed_data = None
+                        if self.parse_filings:
+                            try:
+                                os.makedirs(os.path.dirname(filepath), exist_ok=True)
+                                async with aiofiles.open(filepath, 'wb') as f:
+                                    await f.write(content)
+
+                                parsed_data = parse_sgml_submission(
+                                    content=content.decode(), 
+                                    output_dir=os.path.dirname(filepath) + f'/{url.split("/")[-1].split(".")[0].replace("-", "")}'
+                                )
+                                
+                                try:
+                                    os.remove(filepath)
+                                except Exception as e:
+                                    print(f"\nError deleting original file {filepath}: {str(e)}")
+                                    
+                            except Exception as e:
+                                print(f"\nError parsing {url}: {str(e)}")
+                                try:
+                                    os.remove(filepath)
+                                    parsed_dir = os.path.dirname(filepath) + f'/{url.split("/")[-1].split(".")[0].replace("-", "")}'
+                                    if os.path.exists(parsed_dir):
+                                        import shutil
+                                        shutil.rmtree(parsed_dir)
+                                except Exception as e:
+                                    print(f"\nError cleaning up files for {url}: {str(e)}")
+                        else:
+                            os.makedirs(os.path.dirname(filepath), exist_ok=True)
+                            async with aiofiles.open(filepath, 'wb') as f:
+                                await f.write(content)
+                        
+                        return filepath, parsed_data
+
+                except Exception as e:
+                    print(f"\nError downloading {url}: {str(e)}")
+                    return None
+
+    async def _download_worker(self, pbar):
+        """Worker to process download queue."""
+        while True:
+            try:
+                url, filepath = await self.download_queue.get()
+                result = await self._download_file(url, filepath)
+                if result:
+                    pbar.update(1)
+                self.download_queue.task_done()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                print(f"\nWorker error processing {url}: {str(e)}")
+                self.download_queue.task_done()
+
+    async def _download_and_process(self, urls, output_dir):
+        """Queue-based download processing."""
+        results = []
+        parsed_results = []
+        
+        pbar = tqdm(total=len(urls), desc="Downloading files [Rate: 0/s | 0 MB/s]")
+        self.current_pbar = pbar
+        
+        for url in urls:
+            filename = url.split('/')[-1]
+            filepath = os.path.join(output_dir, filename)
+            await self.download_queue.put((url, filepath))
+        
+        workers = [asyncio.create_task(self._download_worker(pbar)) 
+                  for _ in range(5)]  # Match number of workers to semaphore
+        
+        await self.download_queue.join()
+        
+        for worker in workers:
+            worker.cancel()
+        
+        await asyncio.gather(*workers, return_exceptions=True)
+        
+        pbar.close()
+        self.current_pbar = None
+        return results, parsed_results
+
+    def download_submissions(self, output_dir='filings', cik=None, ticker=None, submission_type=None, date=None, parse=True):
+        """Main method to download SEC filings."""
+        self.parse_filings = parse
+        
+        async def _download():
+            async with self as downloader:
+                if ticker is not None:
+                    cik_value = identifier_to_cik(ticker)
+                else:
+                    cik_value = cik
+
+                params = {}
+                if cik_value:
+                    if isinstance(cik_value, list):
+                        params['ciks'] = ','.join(str(c).zfill(10) for c in cik_value)
+                    else:
+                        params['ciks'] = str(cik_value).zfill(10)
+
+                params['forms'] = ','.join(submission_type) if isinstance(submission_type, list) else submission_type if submission_type else "-0"
+
+                if isinstance(date, list):
+                    dates = [(d, d) for d in date]
+                elif isinstance(date, tuple):
+                    dates = [date]
+                else:
+                    date_str = date if date else f"2001-01-01,{datetime.now().strftime('%Y-%m-%d')}"
+                    start, end = date_str.split(',')
+                    dates = [(start, end)]
+
+                all_filepaths = []
+                all_parsed_data = []
+                
+                for start_date, end_date in dates:
+                    params['startdt'] = start_date
+                    params['enddt'] = end_date
+                    base_url = "https://efts.sec.gov/LATEST/search-index"
+                    efts_url = f"{base_url}?{urlencode(params, doseq=True)}"
+                    
+                    urls = await self._get_filing_urls_from_efts(efts_url)
+                    if urls:
+                        filepaths, parsed_data = await self._download_and_process(urls, output_dir)
+                        all_filepaths.extend(filepaths)
+                        all_parsed_data.extend(parsed_data)
+
+                return all_filepaths, all_parsed_data
+
+        return asyncio.run(_download())
+
+    def download_company_concepts(self, output_dir='company_concepts', cik=None, ticker=None):
+        """Download company concept data."""
+        async def _download_concepts():
+            async with self as downloader:
+                if ticker is not None:
+                    ciks = identifier_to_cik(ticker)
+                elif cik:
+                    ciks = [cik] if not isinstance(cik, list) else cik
+                else:
+                    company_tickers = load_package_csv('company_tickers')
+                    ciks = [company['cik'] for company in company_tickers]
+
+                os.makedirs(output_dir, exist_ok=True)
+                urls = [f'https://data.sec.gov/api/xbrl/companyfacts/CIK{str(cik).zfill(10)}.json' for cik in ciks]
+                
+                pbar = tqdm(total=len(urls), desc="Downloading concepts [Rate: 0/s | 0 MB/s]")
+                self.current_pbar = pbar
+                
+                for url in urls:
+                    filename = url.split('/')[-1]
+                    filepath = os.path.join(output_dir, filename)
+                    await self.download_queue.put((url, filepath))
+                
+                workers = [asyncio.create_task(self._download_worker(pbar)) 
+                          for _ in range(5)]
+                
+                await self.download_queue.join()
+                
+                for worker in workers:
+                    worker.cancel()
+                
+                await asyncio.gather(*workers, return_exceptions=True)
+                
+                pbar.close()
+                self.current_pbar = None
+                
+                results = []
+                for url in urls:
+                    filename = url.split('/')[-1]
+                    filepath = os.path.join(output_dir, filename)
+                    if os.path.exists(filepath):
+                        results.append(filepath)
+                
+                return results
+
+        return asyncio.run(_download_concepts())

datamule/downloader/premiumdownloader.py

@@ -0,0 +1,332 @@
+import os
+import asyncio
+import aiohttp
+from pathlib import Path
+from tqdm import tqdm
+import time
+import shutil
+import ssl
+import zstandard as zstd
+import io
+import json
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from queue import Queue, Empty
+from threading import Thread
+from datamule.parser.sgml_parsing.sgml_parser_cy import parse_sgml_submission
+import urllib.parse
+from ..helper import identifier_to_cik
+
+class InsufficientBalanceError(Exception):
+    def __init__(self, required_cost, current_balance, total_urls):
+        self.required_cost = required_cost
+        self.current_balance = current_balance
+        self.total_urls = total_urls
+        message = (f"Insufficient balance. Required: ${required_cost:.4f}, "
+                  f"Current balance: ${current_balance:.4f}, "
+                  f"Total URLs: {total_urls}")
+        super().__init__(message)
+
+class PremiumDownloader:
+    def __init__(self, api_key=None):
+        self.BASE_URL = "https://library.datamule.xyz/original/nc/"
+        self.API_BASE_URL = "https://sec-library.jgfriedman99.workers.dev/"
+        self.CHUNK_SIZE = 2 * 1024 * 1024
+        self.MAX_CONCURRENT_DOWNLOADS = 100
+        self.MAX_DECOMPRESSION_WORKERS = 16
+        self.MAX_PROCESSING_WORKERS = 16
+        self.QUEUE_SIZE = 10
+        if api_key is not None:
+            self._api_key = api_key
+
+    @property
+    def api_key(self):
+        return getattr(self, '_api_key', None) or os.getenv('DATAMULE_API_KEY')
+
+    @api_key.setter
+    def api_key(self, value):
+        if not value:
+            raise ValueError("API key cannot be empty")
+        self._api_key = value
+
+    def _log_error(self, output_dir, filename, error_msg):
+        error_file = os.path.join(output_dir, 'errors.json')
+        try:
+            if os.path.exists(error_file):
+                with open(error_file, 'r') as f:
+                    errors = json.load(f)
+            else:
+                errors = {}
+            
+            errors[filename] = str(error_msg)
+            
+            with open(error_file, 'w') as f:
+                json.dump(errors, f, indent=2)
+        except Exception as e:
+            print(f"Failed to log error to {error_file}: {str(e)}")
+
+    async def _fetch_submissions(self, session, submission_type=None, cik=None, filing_date=None, page=1):
+        params = {
+            'api_key': self.api_key,
+            'page': page
+        }
+        
+        if submission_type:
+            if isinstance(submission_type, list):
+                params['submission_type'] = ','.join(str(x) for x in submission_type)
+            else:
+                params['submission_type'] = str(submission_type)
+        
+        if cik:
+            if isinstance(cik, list):
+                params['cik'] = ','.join(str(x) for x in cik)
+            else:
+                params['cik'] = str(cik)
+        
+        if filing_date:
+            if isinstance(filing_date, tuple):
+                params['startdt'] = str(filing_date[0])
+                params['enddt'] = str(filing_date[1])
+            else:
+                if isinstance(filing_date, list):
+                    params['filing_date'] = ','.join(str(x) for x in filing_date)
+                else:
+                    params['filing_date'] = str(filing_date)
+
+        url = f"{self.API_BASE_URL}?{urllib.parse.urlencode(params)}"
+        
+        async with session.get(url) as response:
+            data = await response.json()
+            if not data.get('success'):
+                raise ValueError(f"API request failed: {data.get('error')}")
+
+            charges = data['metadata']['billing']['charges']
+            print(f"\nCost: ${charges['results']:.12f} downloads + ${charges['rows_read']:.12f} row reads = ${charges['total']:.12f}")
+            print(f"Balance: ${data['metadata']['billing']['remaining_balance']:.12f}")
+            
+            urls = [f"{self.BASE_URL}{str(sub['accession_number']).zfill(18)}.sgml{'.zst' if sub.get('compressed', '').lower() == 'true' else ''}" for sub in data['data']]
+            return urls, data['metadata']['pagination']
+
+    class FileProcessor:
+        def __init__(self, output_dir, max_workers, queue_size, pbar, downloader):
+            self.processing_queue = Queue(maxsize=queue_size)
+            self.should_stop = False
+            self.processing_workers = []
+            self.output_dir = output_dir
+            self.max_workers = max_workers
+            self.batch_size = 10
+            self.pbar = pbar
+            self.downloader = downloader
+
+        def start_processing_workers(self):
+            for _ in range(self.max_workers):
+                worker = Thread(target=self._processing_worker)
+                worker.daemon = True
+                worker.start()
+                self.processing_workers.append(worker)
+
+        def _process_file(self, item):
+            filename, content = item
+            clean_name = filename[:-4] if filename.endswith('.zst') else filename
+            output_path = os.path.join(self.output_dir, Path(clean_name).stem)
+            try:
+                parse_sgml_submission(None, output_dir=output_path, content=content)
+                self.pbar.update(1)
+            except Exception as e:
+                self.downloader._log_error(self.output_dir, filename, str(e))
+
+        def _processing_worker(self):
+            batch = []
+            while not self.should_stop:
+                try:
+                    item = self.processing_queue.get(timeout=1)
+                    if item is None:
+                        break
+
+                    batch.append(item)
+
+                    if len(batch) >= self.batch_size or self.processing_queue.empty():
+                        for item in batch:
+                            self._process_file(item)
+                            self.processing_queue.task_done()
+                        batch = []
+
+                except Empty:
+                    if batch:
+                        for item in batch:
+                            self._process_file(item)
+                            self.processing_queue.task_done()
+                        batch = []
+
+        def stop_workers(self):
+            self.should_stop = True
+            for _ in self.processing_workers:
+                self.processing_queue.put(None)
+            for worker in self.processing_workers:
+                worker.join()
+
+    def decompress_stream(self, compressed_chunks, filename, output_dir, processor):
+        dctx = zstd.ZstdDecompressor()
+        try:
+            input_buffer = io.BytesIO(b''.join(compressed_chunks))
+            decompressed_content = io.BytesIO()
+            
+            with dctx.stream_reader(input_buffer) as reader:
+                shutil.copyfileobj(reader, decompressed_content)
+                
+            content = decompressed_content.getvalue().decode('utf-8')
+            processor.processing_queue.put((filename, content))
+            return True
+                
+        except Exception as e:
+            self._log_error(output_dir, filename, f"Decompression error: {str(e)}")
+            return False
+        finally:
+            try:
+                input_buffer.close()
+                decompressed_content.close()
+            except:
+                pass
+
+    def save_regular_file(self, chunks, filename, output_dir, processor):
+        try:
+            content = b''.join(chunks).decode('utf-8')
+            processor.processing_queue.put((filename, content))
+            return True
+                
+        except Exception as e:
+            self._log_error(output_dir, filename, f"Error saving file: {str(e)}")
+            return False
+
+    async def download_and_process(self, session, url, semaphore, decompression_pool, output_dir, processor):
+        async with semaphore:
+            chunks = []
+            filename = url.split('/')[-1]
+
+            api_key = self.api_key
+            if not api_key:
+                raise ValueError("No API key found. Please set DATAMULE_API_KEY environment variable or provide api_key in constructor")
+
+            try:
+                headers = {
+                    'Connection': 'keep-alive',
+                    'Accept-Encoding': 'gzip, deflate, br',
+                    'Authorization': f'Bearer {api_key}'
+                }
+                
+                async with session.get(url, headers=headers) as response:
+                    if response.status == 200:
+                        async for chunk in response.content.iter_chunked(self.CHUNK_SIZE):
+                            chunks.append(chunk)
+
+                        loop = asyncio.get_running_loop()
+                        if filename.endswith('.zst'):
+                            success = await loop.run_in_executor(
+                                decompression_pool,
+                                partial(self.decompress_stream, chunks, filename, output_dir, processor)
+                            )
+                        else:
+                            success = await loop.run_in_executor(
+                                decompression_pool,
+                                partial(self.save_regular_file, chunks, filename, output_dir, processor)
+                            )
+
+                        if not success:
+                            self._log_error(output_dir, filename, "Failed to process file")
+                    elif response.status == 401:
+                        self._log_error(output_dir, filename, "Authentication failed: Invalid API key")
+                        raise ValueError("Invalid API key")
+                    else:
+                        self._log_error(output_dir, filename, f"Download failed: Status {response.status}")
+            except Exception as e:
+                self._log_error(output_dir, filename, str(e))
+
+    async def process_batch(self, urls, output_dir):
+        os.makedirs(output_dir, exist_ok=True)
+        
+        with tqdm(total=len(urls), desc="Processing files") as pbar:
+            processor = self.FileProcessor(output_dir, self.MAX_PROCESSING_WORKERS, self.QUEUE_SIZE, pbar, self)
+            processor.start_processing_workers()
+
+            semaphore = asyncio.Semaphore(self.MAX_CONCURRENT_DOWNLOADS)
+            decompression_pool = ThreadPoolExecutor(max_workers=self.MAX_DECOMPRESSION_WORKERS)
+
+            connector = aiohttp.TCPConnector(
+                limit=self.MAX_CONCURRENT_DOWNLOADS,
+                force_close=False,
+                ssl=ssl.create_default_context(),
+                ttl_dns_cache=300,
+                keepalive_timeout=60
+            )
+
+            async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=3600)) as session:
+                tasks = [self.download_and_process(session, url, semaphore, decompression_pool, output_dir, processor) for url in urls]
+                await asyncio.gather(*tasks, return_exceptions=True)
+
+            processor.processing_queue.join()
+            processor.stop_workers()
+            decompression_pool.shutdown()
+
+    async def download_all_pages(self, submission_type=None, cik=None, filing_date=None, output_dir="download"):
+        connector = aiohttp.TCPConnector(ssl=ssl.create_default_context())
+        async with aiohttp.ClientSession(connector=connector) as session:
+            try:
+                urls, pagination = await self._fetch_submissions(session, submission_type=submission_type, cik=cik, filing_date=filing_date, page=1)
+                total_urls = urls.copy()
+                current_page = 1
+                
+                while pagination.get('hasMore', False):
+                    current_page += 1
+                    more_urls, pagination = await self._fetch_submissions(session, submission_type=submission_type, cik=cik, filing_date=filing_date, page=current_page)
+                    total_urls.extend(more_urls)
+                
+                if total_urls:
+                    start_time = time.time()
+                    await self.process_batch(total_urls, output_dir)
+                    elapsed_time = time.time() - start_time
+                    print(f"\nProcessing completed in {elapsed_time:.2f} seconds")
+                else:
+                    print("No submissions found matching the criteria")
+                    
+            except InsufficientBalanceError as e:
+                error_msg = {
+                    "error": "insufficient_balance",
+                    "required_cost": e.required_cost,
+                    "current_balance": e.current_balance,
+                    "total_urls": e.total_urls,
+                    "additional_funds_needed": e.required_cost - e.current_balance
+                }
+                self._log_error(output_dir, "balance_check", error_msg)
+                return
+
+    def download_submissions(self, submission_type=None, cik=None, ticker=None, filing_date=None, output_dir="download"):
+        if self.api_key is None:
+            raise ValueError("No API key found. Please set DATAMULE_API_KEY environment variable or provide api_key in constructor")
+
+        if filing_date is not None:
+            if isinstance(filing_date, str):
+                filing_date = int(filing_date.replace('-', ''))
+            elif isinstance(filing_date, list):
+                filing_date = [int(x.replace('-', '')) for x in filing_date]
+            elif isinstance(filing_date, tuple):
+                filing_date = (int(filing_date[0].replace('-', '')), int(filing_date[1].replace('-', '')))
+
+        if ticker is not None:
+            cik = identifier_to_cik(ticker)
+
+        if cik is not None:
+            if isinstance(cik, str):
+                cik = [int(cik)]
+            elif isinstance(cik, int):
+                cik = [cik]
+            elif isinstance(cik, list):
+                cik = [int(x) for x in cik]
+
+        async def _download():
+            try:
+                await self.download_all_pages(submission_type=submission_type, cik=cik, filing_date=filing_date, output_dir=output_dir)
+            except Exception as e:
+                if not isinstance(e, InsufficientBalanceError):
+                    self._log_error(output_dir, "download_error", str(e))
+
+        asyncio.run(_download())