datamule 0.381__py3-none-any.whl → 1.0.0__py3-none-any.whl

datamule/downloader/premiumdownloader.py

@@ -0,0 +1,335 @@
+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 secsgml 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
+            try:
+                parse_sgml_submission(output_dir=self.output_dir, content=content)
+                self.pbar.update(1)
+            except Exception as e:
+                accession_dir = os.path.join(self.output_dir, filename.split('.')[0])  
+                if os.path.exists(accession_dir):
+                    shutil.rmtree(accession_dir)
+                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
+            )
+
+            # timeout should be max 2 hours.
+            async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=7200)) 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:
+                    total_urls = list(set(total_urls)) # Remove duplicates
+                    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())

datamule/helper.py

@@ -1,136 +1,123 @@
-import requests
-import os
-from tqdm import tqdm
-import zipfile
-from pkg_resources import resource_filename
-import csv
-import re
-
-# Unused in current implementation.
-def construct_primary_doc_url(cik, accession_number,primary_doc_url):
-    accession_number = accession_number.replace("-", "")
-    return f"https://www.sec.gov/Archives/edgar/data/{cik}/{accession_number}/{primary_doc_url}"
-
-# DONE
-def _download_from_dropbox(url, output_path):
-    headers = {'user-agent': 'Wget/1.16 (linux-gnu)'}
-    r = requests.get(url, stream=True, headers=headers)
-    total_size = int(r.headers.get('content-length', 0))
-    
-    with open(output_path, 'wb') as f, tqdm(
-        desc="Downloading " + os.path.basename(output_path),
-        total=total_size,
-        unit='iB',
-        unit_scale=True,
-        unit_divisor=1024,
-    ) as progress_bar:
-        for chunk in r.iter_content(chunk_size=1024):
-            size = f.write(chunk)
-            progress_bar.update(size)
-    
-    # Check if the downloaded file is a zip file
-    if zipfile.is_zipfile(output_path):
-        extract_path = os.path.dirname(output_path)
-        with zipfile.ZipFile(output_path, 'r') as zip_ref:
-            for file_info in zip_ref.infolist():
-                extract_file_path = os.path.join(extract_path, file_info.filename)
-                with zip_ref.open(file_info) as file_in_zip, \
-                    open(extract_file_path, 'wb') as output_file, \
-                    tqdm(total=file_info.file_size, unit='B', unit_scale=True, 
-                         desc=f"Extracting {file_info.filename}") as pbar:
-                    while True:
-                        chunk = file_in_zip.read(8192)
-                        if not chunk:
-                            break
-                        output_file.write(chunk)
-                        pbar.update(len(chunk))
-        
-        # Remove the zip file after extraction
-        os.remove(output_path)
-        print(f"Extracted contents to {extract_path}")
-    else:
-        print(f"Downloaded file is not a zip. Saved to {output_path}")
-
-# May generalize to load any package resource
-def load_package_csv(name):
-    """Load package CSV files"""
-    csv_path = resource_filename('datamule', f'data/{name}.csv')
-    company_tickers = []
-    
-    with open(csv_path, 'r') as csvfile:
-        csv_reader = csv.DictReader(csvfile)
-        for row in csv_reader:
-            company_tickers.append(row)
-    
-    return company_tickers
-
-def load_package_dataset(dataset):
-    if dataset == 'company_tickers':
-        return load_package_csv('company_tickers')
-    elif dataset =='company_former_names':
-        return load_package_csv('company_former_names')
-    elif dataset =='company_metadata':
-        return load_package_csv('company_metadata')
-    elif dataset == 'sec_glossary':
-        return load_package_csv('sec-glossary')
-    elif dataset == 'xbrl_descriptions':
-        return load_package_csv('xbrl_descriptions')
-
-# DONE
-def identifier_to_cik(ticker):
-    """Convert company tickers to CIK codes"""
-    company_tickers = load_package_csv('company_tickers')
-    if ticker:
-        if isinstance(ticker, list):
-            cik = []
-            for t in ticker:
-                cik.extend([company['cik'] for company in company_tickers if t == company['ticker']])
-        else:
-            cik = [company['cik'] for company in company_tickers if ticker == company['ticker']]
-
-    if not cik:
-        raise ValueError("No matching companies found")
-
-    return cik
-
-
-def fix_filing_url(url):
-    """Some Filings URLs have the wrong path. This is an issue with EFTS, that is solved in SEC.gov index page, but SEC.gov still has links to broken url."""
-    # Check if the URL ends with '/0001.txt'
-    if url.endswith('/0001.txt'):
-        # Extract the accession number from the URL
-        match = re.search(r'/(\d{18})/', url)
-        if match:
-            accession_number = match.group(1)
-            # Add dashes to the accession number
-            formatted_accession_number = f"{accession_number[:10]}-{accession_number[10:12]}-{accession_number[12:]}"
-            # Construct the new URL
-            new_url = url.rsplit('/', 1)[0] + f'/{formatted_accession_number}-0001.txt'
-            return new_url
-    elif url.endswith('/0001.htm'):
-        # Extract the accession number from the URL
-        match = re.search(r'/(\d{18})/', url)
-        if match:
-            accession_number = match.group(1)
-            # Add dashes to the accession number
-            formatted_accession_number = f"{accession_number[:10]}-{accession_number[10:12]}-{accession_number[12:]}"
-            # Construct the new URL
-            new_url = url.rsplit('/', 1)[0] + f'/{formatted_accession_number}-0001.htm'
-            return new_url
-    
-    # If the URL doesn't end with '/0001.txt' or doesn't contain a valid accession number,
-    # return the original URL
-    return url
-
-def convert_to_dashed_accession(accession):
-    # Remove any existing dashes or whitespace
-    cleaned = ''.join(accession.split())
-    
-    # Check if the cleaned string has 18 characters
-    if len(cleaned) != 18:
-        raise ValueError("Invalid accession number format. Expected 18 characters.")
-    
-    # Insert dashes at the correct positions
-    dashed = f"{cleaned[:10]}-{cleaned[10:12]}-{cleaned[12:]}"
-    
-    return dashed
+import requests
+import os
+from tqdm import tqdm
+import zipfile
+from pkg_resources import resource_filename
+import csv
+import re
+
+# Unused in current implementation.
+def construct_primary_doc_url(cik, accession_number,primary_doc_url):
+    accession_number = accession_number.replace("-", "")
+    return f"https://www.sec.gov/Archives/edgar/data/{cik}/{accession_number}/{primary_doc_url}"
+
+# DONE
+def _download_from_dropbox(url, output_path):
+    headers = {'user-agent': 'Wget/1.16 (linux-gnu)'}
+    r = requests.get(url, stream=True, headers=headers)
+    total_size = int(r.headers.get('content-length', 0))
+    
+    with open(output_path, 'wb') as f, tqdm(
+        desc="Downloading " + os.path.basename(output_path),
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as progress_bar:
+        for chunk in r.iter_content(chunk_size=1024):
+            size = f.write(chunk)
+            progress_bar.update(size)
+    
+    # Check if the downloaded file is a zip file
+    if zipfile.is_zipfile(output_path):
+        extract_path = os.path.dirname(output_path)
+        with zipfile.ZipFile(output_path, 'r') as zip_ref:
+            for file_info in zip_ref.infolist():
+                extract_file_path = os.path.join(extract_path, file_info.filename)
+                with zip_ref.open(file_info) as file_in_zip, \
+                    open(extract_file_path, 'wb') as output_file, \
+                    tqdm(total=file_info.file_size, unit='B', unit_scale=True, 
+                         desc=f"Extracting {file_info.filename}") as pbar:
+                    while True:
+                        chunk = file_in_zip.read(8192)
+                        if not chunk:
+                            break
+                        output_file.write(chunk)
+                        pbar.update(len(chunk))
+        
+        # Remove the zip file after extraction
+        os.remove(output_path)
+        print(f"Extracted contents to {extract_path}")
+    else:
+        print(f"Downloaded file is not a zip. Saved to {output_path}")
+
+# May generalize to load any package resource
+def load_package_csv(name):
+    """Load package CSV files"""
+    csv_path = resource_filename('datamule', f'data/{name}.csv')
+    company_tickers = []
+    
+    with open(csv_path, 'r') as csvfile:
+        csv_reader = csv.DictReader(csvfile)
+        for row in csv_reader:
+            company_tickers.append(row)
+    
+    return company_tickers
+
+def load_package_dataset(dataset):
+    if dataset == 'company_tickers':
+        return load_package_csv('company_tickers')
+    elif dataset =='company_former_names':
+        return load_package_csv('company_former_names')
+    elif dataset =='company_metadata':
+        return load_package_csv('company_metadata')
+    elif dataset == 'sec_glossary':
+        return load_package_csv('sec-glossary')
+    elif dataset == 'xbrl_descriptions':
+        return load_package_csv('xbrl_descriptions')
+
+# DONE
+def identifier_to_cik(ticker):
+    """Convert company tickers to CIK codes"""
+    company_tickers = load_package_csv('company_tickers')
+    if ticker:
+        if isinstance(ticker, list):
+            cik = []
+            for t in ticker:
+                cik.extend([company['cik'] for company in company_tickers if t == company['ticker']])
+        else:
+            cik = [company['cik'] for company in company_tickers if ticker == company['ticker']]
+
+    if not cik:
+        raise ValueError("No matching companies found")
+
+    return cik
+
+
+def fix_filing_url(url):
+    match_suffix = re.search(r'/(\d{4})\.(.+?)$', url)
+    if match_suffix:
+        suffix_number = match_suffix.group(1)
+        file_ext = match_suffix.group(2)
+        match_accession = re.search(r'/(\d{18})/', url)
+        if match_accession:
+            accession_number = match_accession.group(1)
+            formatted_accession_number = f"{accession_number[:10]}-{accession_number[10:12]}-{accession_number[12:]}"
+            new_url = url.rsplit('/', 1)[0] + f'/{formatted_accession_number}-{suffix_number}.{file_ext}'
+            return new_url
+    return url
+
+def convert_to_dashed_accession(accession):
+    # Remove any existing dashes or whitespace
+    cleaned = ''.join(accession.split())
+    
+    # Check if the cleaned string has 18 characters
+    if len(cleaned) != 18:
+        raise ValueError("Invalid accession number format. Expected 18 characters.")
+    
+    # Insert dashes at the correct positions
+    dashed = f"{cleaned[:10]}-{cleaned[10:12]}-{cleaned[12:]}"
+    
+    return dashed
+
+headers = {'User-Agent': 'John Smith johnsmith@gmail.com'}