datamule 1.4.9__py3-none-any.whl → 1.5.1__py3-none-any.whl

datamule/portfolio.py

@@ -33,7 +33,8 @@ class Portfolio:
         self.api_key = api_key
     
     def _load_submissions(self):
-        folders = [f for f in self.path.iterdir() if f.is_dir()]
+        folders = [f for f in self.path.iterdir() if f.is_dir() or f.suffix=='.tar']
+        print(folders)
         print(f"Loading {len(folders)} submissions")
         
         def load_submission(folder):
@@ -125,7 +126,7 @@ class Portfolio:
             # First query, just set the accession numbers
             self.accession_numbers = new_accession_numbers
 
-    def download_submissions(self, cik=None, ticker=None, submission_type=None, filing_date=None, provider=None,document_type=None,requests_per_second=5, **kwargs):
+    def download_submissions(self, cik=None, ticker=None, submission_type=None, filing_date=None, provider=None,document_type=[],requests_per_second=5, **kwargs):
         if provider is None:
             config = Config()
             provider = config.get_default_source()

datamule/sec/submissions/downloader.py

@@ -1,55 +1,23 @@
 import os
-import json
 from .streamer import stream
-import aiofiles
-from ...submission import Submission
-
-async def download_callback(hit, content, cik, accno, url, output_dir="filings", keep_document_types=None):
-    """Save downloaded SEC submission to disk."""
-    try:
-        # Create a Submission object directly from the content
-        # Note: the content needs to be decoded from bytes to string for the parser
-        submission = Submission(sgml_content=content, 
-                               keep_document_types=keep_document_types)
-        
-        # Use the async save method to write the submission to disk
-        file_dir = await submission.save_async(output_dir=output_dir)
-        
-        return file_dir
-    except Exception as e:
-        print(f"Error processing {accno}: {e}")
-        return None
+from secsgml import write_sgml_file_to_tar
+from tqdm import tqdm
 
 def download(cik=None, submission_type=None, filing_date=None, location=None, name=None, 
              requests_per_second=5, output_dir="filings", accession_numbers=None, 
-             quiet=False, keep_document_types=None):
-    """
-    Download SEC EDGAR filings and extract their documents.
-    
-    Parameters:
-    - cik: CIK number(s) to query for
-    - submission_type: Filing type(s) to query for (default: 10-K)
-    - filing_date: Date or date range to query for
-    - location: Location code to filter by (e.g., 'CA' for California)
-    - name: Company name to search for (alternative to providing CIK)
-    - requests_per_second: Rate limit for SEC requests
-    - output_dir: Directory to save documents
-    - accession_numbers: Optional list of accession numbers to filter by
-    - quiet: Whether to suppress progress output
-    - keep_document_types: Optional list of document types to keep (e.g. ['10-K', 'EX-10.1'])
-    
-    Returns:
-    - List of all document paths processed
-    """
+             quiet=False, keep_document_types=[]):
     # Make sure output directory exists
     os.makedirs(output_dir, exist_ok=True)
-    
+
+    pbar = tqdm(desc="Writing", unit=" submissions", disable=quiet,position=2)
+
     # Create a wrapper for the download_callback that includes the output_dir
     async def callback_wrapper(hit, content, cik, accno, url):
-        return await download_callback(hit, content, cik, accno, url, 
-                                     output_dir=output_dir,
-                                     keep_document_types=keep_document_types)
-    
+        output_path = os.path.join(output_dir, accno.replace('-','') + '.tar')
+        write_sgml_file_to_tar(output_path, bytes_content=content, filter_document_types=keep_document_types)
+        pbar.update(1)
+
+
     # Call the stream function with our callback
     return stream(
         cik=cik,

datamule/seclibrary/downloader.py

@@ -15,6 +15,7 @@ from threading import Thread
 from .query import query
 from os import cpu_count
 from ..submission import Submission
+from secsgml import write_sgml_file_to_tar
 
 
 
@@ -73,7 +74,7 @@ class Downloader:
             print(f"Failed to log error to {error_file}: {str(e)}")
 
     class FileProcessor:
-        def __init__(self, output_dir, max_workers, queue_size, pbar, downloader, keep_document_types=None):
+        def __init__(self, output_dir, max_workers, queue_size, pbar, downloader, keep_document_types=[]):
             self.processing_queue = Queue(maxsize=queue_size)
             self.should_stop = False
             self.processing_workers = []
@@ -93,17 +94,9 @@ class Downloader:
 
         def _process_file(self, item):
             filename, content = item
-            try:
-                submission = Submission(sgml_content=content, keep_document_types=self.keep_document_types)
-                # Use the shared event loop to run save_async
-                self.downloader._run_coroutine(submission.save_async(output_dir=self.output_dir))
-                self.pbar.update(1)
-            except Exception as e:
-                print(f"Exception {e} in {filename}")
-                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))
+            output_path = os.path.join(self.output_dir, filename.split('.')[0] + '.tar')
+            write_sgml_file_to_tar(output_path, bytes_content=content, filter_document_types=self.keep_document_types)
+            self.pbar.update(1)
 
         def _processing_worker(self):
             batch = []
@@ -211,7 +204,7 @@ class Downloader:
             except Exception as e:
                 self._log_error(output_dir, filename, str(e))
 
-    async def process_batch(self, urls, output_dir, keep_document_types=None):
+    async def process_batch(self, urls, output_dir, keep_document_types=[]):
         os.makedirs(output_dir, exist_ok=True)
         
         with tqdm(total=len(urls), desc="Processing files") as pbar:
@@ -238,7 +231,7 @@ class Downloader:
             processor.stop_workers()
             decompression_pool.shutdown()
 
-    def download(self, submission_type=None, cik=None, filing_date=None, output_dir="downloads", accession_numbers=None, keep_document_types=None):
+    def download(self, submission_type=None, cik=None, filing_date=None, output_dir="downloads", accession_numbers=None, keep_document_types=[]):
         """
         Query SEC filings and download/process them.
         
@@ -299,7 +292,7 @@ class Downloader:
             self.loop.call_soon_threadsafe(self.loop.stop)
 
 
-def download(submission_type=None, cik=None, filing_date=None, api_key=None, output_dir="downloads", accession_numbers=None, keep_document_types=None):
+def download(submission_type=None, cik=None, filing_date=None, api_key=None, output_dir="downloads", accession_numbers=None, keep_document_types=[]):
     """
     Query SEC filings and download/process them.
     

datamule/submission.py

@@ -2,10 +2,12 @@ from pathlib import Path
 import json
 from .document.document import Document
 from secsgml import parse_sgml_content_into_memory
-import os
-import aiofiles
+import tarfile
+import shutil
+import zstandard as zstd
+from io import BytesIO
+import gzip
 
-# TODO add .tar path
 class Submission:
     def __init__(self, path=None,sgml_content=None,keep_document_types=None):
         if path is None and sgml_content is None:
@@ -25,7 +27,7 @@ class Submission:
             filtered_metadata_documents = []
 
             for idx,doc in enumerate(self.metadata.content['documents']):
-                type = doc.get('type').upper()
+                type = doc.get('type')()
                 
                 # Keep only specified types
                 if keep_document_types is not None and type not in keep_document_types:
@@ -34,7 +36,7 @@ class Submission:
                 # write as txt if not declared
                 filename = doc.get('filename','.txt')
                 extension = Path(filename).suffix
-                self.documents.append(Document(type=type.upper(), content=raw_documents[idx], extension=extension,filing_date=self.filing_date,accession=self.accession))
+                self.documents.append(Document(type=type, content=raw_documents[idx], extension=extension,filing_date=self.filing_date,accession=self.accession))
 
                 filtered_metadata_documents.append(doc)
             
@@ -42,156 +44,172 @@ class Submission:
 
         if path is not None:
             self.path = Path(path)  
-            metadata_path = self.path / 'metadata.json'
-            with metadata_path.open('r') as f:
-                metadata = json.load(f) 
-            self.metadata = Document(type='submission_metadata', content=metadata, extension='.json',filing_date=None,accession=None,path=metadata_path)
+            if self.path.suffix == '.tar':
+                with tarfile.open(self.path,'r') as tar:
+                    metadata_obj = tar.extractfile('metadata.json')
+                    metadata = json.loads(metadata_obj.read().decode('utf-8'))
 
-            # Code dupe
+                # tarpath
+                metadata_path = f"{self.path}::metadata.json"
+            else:
+                metadata_path = self.path / 'metadata.json'
+                with metadata_path.open('r') as f:
+                    metadata = json.load(f) 
+            self.metadata = Document(type='submission_metadata', content=metadata, extension='.json',filing_date=None,accession=None,path=metadata_path)
             self.accession = self.metadata.content['accession-number']
             self.filing_date= f"{self.metadata.content['filing-date'][:4]}-{self.metadata.content['filing-date'][4:6]}-{self.metadata.content['filing-date'][6:8]}"
     
 
 
-
-    def document_type(self, document_type):
-        # Convert single document type to list for consistent handling
-        if isinstance(document_type, str):
-            document_types = [document_type.lower()]
-        else:
-            document_types = [item.lower() for item in document_type]
-
-        for idx,doc in enumerate(self.metadata.content['documents']):
-            if doc['type'] in document_types:
-                
-                # if loaded from path
-                if self.path is not None:
-                    filename = doc.get('filename')
-                    # oh we need handling here for sequences case
-                    if filename is None:
-                        filename = doc['sequence'] + '.txt'
-                        
-                    document_path = self.path / filename
-                    extension = document_path.suffix
-
-                    with document_path.open('rb') as f:
-                        content = f.read()
-
-                    if extension in ['.htm','.html','.txt','.xml']:
-                        content = content.decode('utf-8', errors='replace')
-
-                    yield Document(type=doc['type'].upper(), content=content, extension=extension,filing_date=self.filing_date,accession=self.accession,path=document_path)
-                # if loaded from sgml_content
-                else:
-                    yield self.documents[idx]
-
-    
-    def __iter__(self):
-        for idx,doc in enumerate(self.metadata.content['documents']):
-            # if loaded from path
-            if self.path is not None:
+    def compress(self, compression=None, level=None, threshold=1048576):
+        if self.path is None:
+            raise ValueError("Compress requires path")
+        
+        if compression is not None and compression not in ['gzip', 'zstd']:
+            raise ValueError("compression must be 'gzip' or 'zstd'")
+        
+        # Create tar file (replace directory with .tar file)
+        tar_path = self.path.with_suffix('.tar')
+        
+        with tarfile.open(tar_path, 'w') as tar:
+            # Add metadata.json first
+            metadata_path = self.path / 'metadata.json'
+            if metadata_path.exists():
+                tar.add(metadata_path, arcname='metadata.json')
+            
+            # Add documents in order
+            for doc in self.metadata.content['documents']:
                 filename = doc.get('filename')
-
-                # oh we need handling here for sequences case
                 if filename is None:
                     filename = doc['sequence'] + '.txt'
-                    
-                document_path = self.path / filename
-                extension = document_path.suffix
-
-                # check if the file exists
-                if document_path.exists():
-                    with document_path.open('rb') as f:
-                        content = f.read()
-
-                    if extension in ['.htm','.html','.txt','.xml']:
-                        content = content.decode('utf-8', errors='replace')
-
-                    yield Document(type=doc['type'].upper(), content=content, extension=extension,filing_date=self.filing_date,accession=self.accession,path=document_path)
-                else:
-                    print(f"Warning: File {document_path} does not exist likely due to keep types in downloading.")
-
-            # if loaded from sgml_content
-            else:
-                yield self.documents[idx]
-
-
-
+                
+                file_path = self.path / filename
+                if file_path.exists():
+                    file_size = file_path.stat().st_size
 
-    def save(self, output_dir="filings"):
-        file_dir = Path(output_dir) / str(self.accession)
-        file_dir.mkdir(parents=True, exist_ok=True)
-        
-        metadata_path = file_dir / "metadata.json"
-        with open(metadata_path, 'w') as f:
-            json.dump(self.metadata.content, f, indent=4)
-        
-        for idx, doc in enumerate(self.metadata.content['documents']):
-            filename = doc.get('filename')
-            if filename is None:
-                filename = f"{doc.get('sequence')}.txt"
-            
-            doc_path = file_dir / filename
-            
-            if self.path is not None:
-                if hasattr(self, 'documents') and self.documents:
-                    content = self.documents[idx].content
-                else:
-                    orig_doc_path = self.path / filename
-                    if orig_doc_path.exists():
-                        with open(orig_doc_path, 'r', encoding='utf-8', errors='replace') as f:
-                            content = f.read()
+                    
+                    # Compress if compression specified and over threshold
+                    if compression is not None and file_size >= threshold:
+                        content = file_path.read_bytes()
+                        
+                        if compression == 'gzip':
+                            compressed_content = gzip.compress(content, compresslevel=level or 6)
+                            compressed_filename = filename + '.gz'
+                        else:  # zstd
+                            cctx = zstd.ZstdCompressor(level=level or 3)
+                            compressed_content = cctx.compress(content)
+                            compressed_filename = filename + '.zst'
+                        
+                        # Add compressed file to tar
+                        tarinfo = tarfile.TarInfo(name=compressed_filename)
+                        tarinfo.size = len(compressed_content)
+                        tar.addfile(tarinfo, BytesIO(compressed_content))
                     else:
-                        print(f"Warning: File {orig_doc_path} does not exist, skipping.")
-                        continue
-            else:
-                content = self.documents[idx].content
-            
-            if isinstance(content, bytes):
-                with open(doc_path, 'wb') as f:
-                    f.write(content)
-            else:
-                with open(doc_path, 'w', encoding='utf-8', errors='replace') as f:
-                    f.write(content)
+                        # Add uncompressed file
+                        tar.add(file_path, arcname=filename)
         
-        return file_dir
-
-    async def save_async(self, output_dir="filings"):
-        file_dir = Path(output_dir) / str(self.accession)
-        os.makedirs(file_dir, exist_ok=True)
+        # Delete original folder
+        shutil.rmtree(self.path)
         
-        metadata_path = file_dir / "metadata.json"
-        async with aiofiles.open(metadata_path, 'w') as f:
-            await f.write(json.dumps(self.metadata.content, indent=4))
+        # Update path to point to new tar file
+        self.path = tar_path
+
+    def decompress(self):
+        if self.path is None:
+            raise ValueError("Decompress requires path")
+        elif self.path.suffix != '.tar':
+            raise ValueError("Can only decompress tar")
         
-        for idx, doc in enumerate(self.metadata.content['documents']):
-            filename = doc.get('filename')
-            # oh we need handling here for sequences case
-            if filename is None:
-                filename = doc['sequence'] + '.txt'
-
-            
-            doc_path = file_dir / filename
-            
-            if self.path is not None:
-                if hasattr(self, 'documents') and self.documents:
-                    content = self.documents[idx].content
-                else:
-                    orig_doc_path = self.path / filename
-                    if orig_doc_path.exists():
-                        async with aiofiles.open(orig_doc_path, 'r', encoding='utf-8', errors='replace') as f:
-                            content = await f.read()
+        # Create output directory (path without .tar extension)
+        output_dir = self.path.with_suffix('')
+        output_dir.mkdir(exist_ok=True)
+        
+        with tarfile.open(self.path, 'r') as tar:
+            for member in tar.getmembers():
+                if member.isfile():
+                    content = tar.extractfile(member).read()
+                    
+                    # Decompress if gzipped
+                    if member.name.endswith('.gz'):
+                        content = gzip.decompress(content)
+                        output_path = output_dir / member.name[:-3]  # Remove .gz extension
                     else:
-                        print(f"Warning: File {orig_doc_path} does not exist, skipping.")
-                        continue
-            else:
-                content = self.documents[idx].content
-            
-            if isinstance(content, bytes):
-                async with aiofiles.open(doc_path, 'wb') as f:
-                    await f.write(content)
-            else:
-                async with aiofiles.open(doc_path, 'w', encoding='utf-8', errors='replace') as f:
-                    await f.write(content)
+                        output_path = output_dir / member.name
+                    
+                    # Write to output directory
+                    output_path.parent.mkdir(parents=True, exist_ok=True)
+                    with output_path.open('wb') as f:
+                        f.write(content)
+
+        # delete original file
+        self.path.unlink()
+        self.path = output_dir
+
+    def _load_document_by_index(self, idx):
+        """Load a document by its index in the metadata documents list."""
+        doc = self.metadata.content['documents'][idx]
         
-        return file_dir
+        # If loaded from sgml_content, return pre-loaded document
+        if self.path is None:
+            return self.documents[idx]
+        
+        # If loaded from path, load document on-demand
+        filename = doc.get('filename')
+        if filename is None:
+            filename = doc['sequence'] + '.txt'
+
+        document_path = self.path / filename
+        extension = document_path.suffix
+
+        if self.path.suffix == '.tar':
+            with tarfile.open(self.path, 'r') as tar:
+                # bandaid fix TODO
+                try:
+                    content = tar.extractfile(filename).read()
+                except:
+                    try:
+                        content = tar.extractfile(filename+'.gz').read()
+                    except:
+                        try: 
+                            content = tar.extractfile(filename+'.zst').read()
+                        except:
+                            raise ValueError("Something went wrong with tar")
+                # Decompress if compressed
+                if filename.endswith('.gz'):
+                    content = gzip.decompress(content)
+                elif filename.endswith('.zst'):
+                    dctx = zstd.ZstdDecompressor()
+                    content = dctx.decompress(content)
+        else:
+            with document_path.open('rb') as f:
+                content = f.read()
+
+            # Decode text files
+            if extension in ['.htm', '.html', '.txt', '.xml']:
+                content = content.decode('utf-8', errors='replace')
+
+        return Document(
+            type=doc['type'], 
+            content=content, 
+            extension=extension,
+            filing_date=self.filing_date,
+            accession=self.accession,
+            path=document_path
+        )
+
+    def __iter__(self):
+        """Make Submission iterable by yielding all documents."""
+        for idx in range(len(self.metadata.content['documents'])):
+            yield self._load_document_by_index(idx)
+
+    def document_type(self, document_type):
+        """Yield documents matching the specified type(s)."""
+        # Convert single document type to list for consistent handling
+        if isinstance(document_type, str):
+            document_types = [document_type]
+        else:
+            document_types = [item for item in document_type]
+
+        for idx, doc in enumerate(self.metadata.content['documents']):
+            if doc['type'] in document_types:
+                yield self._load_document_by_index(idx)

{datamule-1.4.9.dist-info → datamule-1.5.1.dist-info}/METADATA

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

{datamule-1.4.9.dist-info → datamule-1.5.1.dist-info}/RECORD

@@ -3,9 +3,9 @@ datamule/config.py,sha256=Y--CVv7JcgrjJkMOSLrvm2S8B9ost6RMSkGviP-MKtg,883
 datamule/helper.py,sha256=KqhAmTMdvATEh3I-O4xLcAcrHB9zXQERBuwzue7zyQw,3674
 datamule/index.py,sha256=Rrcna9FJV-Oh_K6O2IuUEIDmtay_7UZ4l4jgKCi7A7I,2079
 datamule/package_updater.py,sha256=Z9zaa_y0Z5cknpRn8oPea3gg4kquFHfpfhduKKCZ6NU,958
-datamule/portfolio.py,sha256=8fiK-vfZM5-NJSvOEsDR2YDb-2njjzFk6l7BiRyrzOM,7168
+datamule/portfolio.py,sha256=PEpi9Sg6aldgafDeIiXVe47Nm9RkhDtwA9ykUDGd6T0,7209
 datamule/sheet.py,sha256=TvFqK9eAYuVoJ2uWdAlx5EN6vS9lke-aZf7FqtUiDBc,22304
-datamule/submission.py,sha256=UTGIkXOFVrBTHLJxTekw60Nvp92GuSsgQApJWXHNuNg,8493
+datamule/submission.py,sha256=Pq6dGf1zjlFdOIKaaNvPgYGkqcqa1PEla-3QcSGqMSk,9300
 datamule/data/listed_filer_metadata.csv,sha256=dT9fQ8AC5P1-Udf_UF0ZkdXJ88jNxJb_tuhi5YYL1rc,2426827
 datamule/document/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datamule/document/document.py,sha256=VaJWo9HrcODlbifYcXzifW3xBD7nUOWAN8zcVCDWMcs,13958
@@ -44,7 +44,7 @@ datamule/sec/utils.py,sha256=JUxwijJiqRMnRJNQzVUamyF5h9ZGc7RnO_zsLOIM73g,2079
 datamule/sec/infrastructure/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datamule/sec/infrastructure/submissions_metadata.py,sha256=f1KarzFSryKm0EV8DCDNsBw5Jv0Tx5aljiGUJkk7DRk,18745
 datamule/sec/submissions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datamule/sec/submissions/downloader.py,sha256=izaz559PtBCAWPWGzqUReloawJtXwnraclgXdzEOteI,2631
+datamule/sec/submissions/downloader.py,sha256=QQ7DkSaXBe-TQ2oZb2coGn3TBXIHywXn0uCGeIaVWCI,1216
 datamule/sec/submissions/eftsquery.py,sha256=mSZon8rlW8dxma7M49ZW5V02Fn-ENOdt9TNO6elBrhE,27983
 datamule/sec/submissions/monitor.py,sha256=dZYuVCi_X82eYA8l_9cbnkRjiawz3K4U-FnCAyJcgk4,7892
 datamule/sec/submissions/streamer.py,sha256=EXyWNCD9N6mZmvm9lFSCFodF19zSQ8jfIbWPZNp0K5Y,11253
@@ -56,9 +56,9 @@ datamule/sec/xbrl/streamcompanyfacts.py,sha256=Qq88PqW5_j1k3Aqrl0KRmKeF54D6Wbb6H
 datamule/sec/xbrl/xbrlmonitor.py,sha256=TKFVfSyyUUfUgFQw4WxEVs4g8Nh-2C0tygNIRmTqW3Y,5848
 datamule/seclibrary/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datamule/seclibrary/bq.py,sha256=C8sb_rpXTvchprrFLcbRar4Qi0XWW25tnv1YsHSS5o4,18025
-datamule/seclibrary/downloader.py,sha256=VIdaQq5wDcYWnqrv9t8J7z0KtdNRGK8ahfBsgvTfdQQ,13675
+datamule/seclibrary/downloader.py,sha256=m8yrRjjRVJ2WtKIKIO_66lQHtAsWssyTPip6B2mW9tE,13275
 datamule/seclibrary/query.py,sha256=qGuursTERRbOGfoDcYcpo4oWkW3PCBW6x1Qf1Puiak4,7352
-datamule-1.4.9.dist-info/METADATA,sha256=ojII-iQoWPNIJXs28FNsQjExKxOdQROgYsIHdOfMw6I,469
-datamule-1.4.9.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-datamule-1.4.9.dist-info/top_level.txt,sha256=iOfgmtSMFVyr7JGl_bYSTDry79JbmsG4p8zKq89ktKk,9
-datamule-1.4.9.dist-info/RECORD,,
+datamule-1.5.1.dist-info/METADATA,sha256=YB8k9NEeb_P77ryO5EGZJiG4wxi6dU2mTYJ0F4Ex--k,469
+datamule-1.5.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+datamule-1.5.1.dist-info/top_level.txt,sha256=iOfgmtSMFVyr7JGl_bYSTDry79JbmsG4p8zKq89ktKk,9
+datamule-1.5.1.dist-info/RECORD,,