sovai 0.2.26__tar.gz → 0.2.27__tar.gz

{sovai-0.2.26 → sovai-0.2.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sovai
-Version: 0.2.26
+Version: 0.2.27
 Summary: python for asset management
 License: Proprietary License
 Author: Derek Snow

{sovai-0.2.26 → sovai-0.2.27}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sovai"
-version = "0.2.26"
+version = "0.2.27"
 description = "python for asset management"
 authors = ["Derek Snow <d.snow@outlook.com>"]
 license = "Proprietary License"
@@ -40,6 +40,8 @@ python-dotenv = ">=0.10"
 polars = ">=0.20.30"
 boto3 = ">=1.20"
 joblib = ">=1.0"
+rapidfuzz = "^3.9.7"
+edgar-tool = "^1.3.3"
 scipy = {version = ">=1.0", optional = true}
 matplotlib = {version = ">=3.0", optional = true}
 plotly = {version = ">=5.0", optional = true}
@@ -59,8 +61,7 @@ lightgbm = {version = ">=4.5.0", optional = true}
 ipywidgets = {version = ">=8.1.3", optional = true}
 polars-talib = {version = "0.1.3", optional = true}
 dash-bootstrap-components = {version = ">=1.6.0", optional = true}
-rapidfuzz = "^3.9.7"
-edgar-tool = "^1.3.3"
+
 
 [tool.poetry.extras]
 full = [
@@ -85,6 +86,17 @@ full = [
  "dash-bootstrap-components"
 ]
 
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+jinja2 = "^3.1.4"
+traitlets = "^5.14.3"
+jupyterlab = "^4.2.5"
+pytest = "^7.0"
+black = "^22.0"
+
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

{sovai-0.2.26 → sovai-0.2.27}/sovai/__init__.py

@@ -6,8 +6,10 @@ from .get_data import data
 from .basic_auth import basic_auth
 from .token_auth import token_auth
 from .utils.file_management import save_or_update_tickers, save_or_update_codes
+
+
 import warnings
-warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore")
 
 
 __version__ = "0.1.25"

{sovai-0.2.26 → sovai-0.2.27}/sovai/get_data.py

@@ -21,6 +21,7 @@ from sovai.utils.converter import convert_data2df
 from sovai.utils.stream import stream_data, stream_data_pyarrow
 from sovai.utils.datetime_formats import datetime_format
 from sovai.utils.client_side import client_side_frame
+from sovai.utils.client_side_s3 import load_frame_s3
 
 ## Note run these for TEST - IT COULD BE HERE some silent problems
 try:
@@ -281,7 +282,7 @@ def normalize_endpoint(endpoint):
     return endpoint.strip("/").strip()
 
 
-client_side_endpoints = {
+client_side_endpoints_gcs = {
     "ratios/relative",
     "market/prices",
     "market/closeadj",
@@ -291,6 +292,10 @@ client_side_endpoints = {
     "lobbying/public",
 }
 
+client_side_endpoints_s3 = {
+    "sec/10k"
+}
+
 # In your main function or wherever this logic is implemented:
 
 
@@ -577,16 +582,25 @@ def data(
     normalized_endpoint = normalize_endpoint(endpoint)
 
     if (
-        normalized_endpoint in client_side_endpoints
+        normalized_endpoint in client_side_endpoints_gcs
         and tickers is not None
         and frequency is None
     ):
-        verbose_mode.vprint("Grabbing client side")
+        verbose_mode.vprint("Grabbing gcs client side")
         _query_cache[cache_key] = client_side_frame(
             normalized_endpoint, tickers, columns, start_date, end_date
         )
         return _query_cache[cache_key]
 
+    elif (
+        normalized_endpoint in client_side_endpoints_s3
+        and tickers is not None
+        and frequency is None
+    ):
+        verbose_mode.vprint("Grabbing s3 client side")
+        _query_cache[cache_key] = load_frame_s3(
+            normalized_endpoint, tickers, columns, start_date, end_date)
+        return _query_cache[cache_key]
     try:
         if tickers is not None and not is_all(tickers):
             params = ticker_mapper(params,verbose)

sovai-0.2.27/sovai/tools/authentication/authentication.py

@@ -0,0 +1,98 @@
+import sovai as sov
+from sovai.api_config import ApiConfig
+import requests
+import pickle
+import base64
+from pyarrow.fs import S3FileSystem
+
+def get_s3_filesystem_pickle(storage_provider: str, verbose: bool = False):
+    """
+    Retrieve the S3FileSystem for the specified storage provider using pickle.
+    
+    :param storage_provider: Either 'wasabi' or 'digitalocean'
+    :param verbose: If True, print additional information about the request
+    :return: S3FileSystem object, or None if the request fails
+    """
+    if storage_provider.lower() not in ['wasabi', 'digitalocean']:
+        raise ValueError("storage_provider must be either 'wasabi' or 'digitalocean'")
+
+    url = f"{ApiConfig.base_url}/credentials/get_s3_filesystem_pickle"
+    headers = {
+        "Authorization": f"Bearer {ApiConfig.token}",
+        "Content-Type": "application/json"
+    }
+    params = {"storage_provider": storage_provider.lower()}
+
+    try:
+        response = requests.get(url, headers=headers, params=params)
+        response.raise_for_status()
+        
+        encoded = response.content
+        pickled = base64.b64decode(encoded)
+        s3fs = pickle.loads(pickled)
+
+        return s3fs
+
+    except requests.RequestException as e:
+        if verbose:
+            print(f"Failed to retrieve S3FileSystem for {storage_provider}: {e}")
+        return None
+
+def get_s3_filesystem_json(storage_provider: str, verbose: bool = False):
+    """
+    Retrieve the S3FileSystem for the specified storage provider using JSON.
+    
+    :param storage_provider: Either 'wasabi' or 'digitalocean'
+    :param verbose: If True, print additional information about the request
+    :return: S3FileSystem object, or None if the request fails
+    """
+    if storage_provider.lower() not in ['wasabi', 'digitalocean']:
+        raise ValueError("storage_provider must be either 'wasabi' or 'digitalocean'")
+
+    url = f"{ApiConfig.base_url}/credentials/get_s3_filesystem_json"
+    headers = {
+        "Authorization": f"Bearer {ApiConfig.token}",
+        "Content-Type": "application/json"
+    }
+    params = {"storage_provider": storage_provider.lower()}
+
+    try:
+        response = requests.get(url, headers=headers, params=params)
+        response.raise_for_status()
+        
+        cred = response.json()
+
+        if verbose:
+            print(f"Successfully retrieved S3 credentials for {storage_provider}")
+
+        s3fs = S3FileSystem(
+            access_key=cred['access_key'],
+            secret_key=cred['secret_key'],
+            endpoint_override=cred['endpoint']
+        )
+
+        try:
+            s3fs.ls('/')
+            if verbose:
+                print(f"Successfully verified S3FileSystem for {storage_provider}")
+            return s3fs
+        except Exception as e:
+            if verbose:
+                print(f"Failed to verify S3FileSystem: {str(e)}")
+            return None
+
+    except requests.RequestException as e:
+        if verbose:
+            print(f"Failed to retrieve S3 credentials for {storage_provider}: {e}")
+        return None
+
+# Usage example:
+# s3fs = get_s3_filesystem_pickle('wasabi', verbose=True)
+# or
+# s3fs = get_s3_filesystem_json('wasabi', verbose=True)
+# if s3fs:
+#     # Use s3fs in your operations
+#     files = s3fs.ls('/')
+#     print(files)
+# else:
+#     print("Failed to get S3FileSystem")

sovai-0.2.27/sovai/utils/client_side_s3.py

@@ -0,0 +1,212 @@
+import pyarrow.dataset as ds
+from pyarrow.fs import S3FileSystem
+import pyarrow as pa
+import pandas as pd
+from functools import lru_cache
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from sovai.tools.authentication import authentication
+
+
+# Try to import CustomDataFrame, use regular DataFrame if not available
+try:
+    from sovai.extensions.pandas_extensions import CustomDataFrame
+    HAS_CUSTOM_DATAFRAME = True
+except ImportError:
+    HAS_CUSTOM_DATAFRAME = False
+    CustomDataFrame = pd.DataFrame  # Fallback to regular DataFrame
+
+
+@lru_cache(maxsize=2)
+def get_cached_s3_filesystem(storage_provider):
+    return authentication.get_s3_filesystem_pickle(storage_provider, verbose=True)
+
+def process_partition(storage_provider, base_path, identifier_column, identifier=None, columns=None, filters=None):
+    s3 = get_cached_s3_filesystem(storage_provider)
+    
+    if identifier:
+        base_path += f"/{identifier_column}={identifier}"
+    
+    dataset = ds.dataset(base_path, filesystem=s3, format='parquet')
+    
+    if filters:
+        operator_map = {
+            '==': lambda a, b: a == b,
+            '!=': lambda a, b: a != b,
+            '<': lambda a, b: a < b,
+            '<=': lambda a, b: a <= b,
+            '>': lambda a, b: a > b,
+            '>=': lambda a, b: a >= b,
+            'in': lambda a, b: a.isin(b),
+            'not in': lambda a, b: ~a.isin(b),
+            'like': lambda a, b: a.match_substring(b),
+        }
+        schema = dataset.schema
+        filter_expr = None
+        for col, op, val in filters:
+            if op not in operator_map:
+                raise ValueError(f"Unsupported operator '{op}' in filters.")
+            field = ds.field(col)
+            field_type = schema.field(col).type
+            if pa.types.is_timestamp(field_type) or pa.types.is_date(field_type):
+                val = pd.to_datetime(val)
+            elif pa.types.is_integer(field_type):
+                val = int(val)
+            elif pa.types.is_floating(field_type):
+                val = float(val)
+            condition = operator_map[op](field, val)
+            filter_expr = condition if filter_expr is None else filter_expr & condition
+    else:
+        filter_expr = None
+    
+    table = dataset.to_table(columns=columns, filter=filter_expr, use_threads=True)
+    df = table.to_pandas(use_threads=True)
+    
+    if identifier:
+        df[identifier_column] = identifier
+    
+    return df
+
+import numpy as np
+
+def client_side_s3_frame(config, identifiers, columns, start_date, end_date):
+    storage_provider = "digitalocean"  # Default to DigitalOcean
+    base_path = config["storage_provider"][storage_provider]
+    identifier_column = config.get("identifier_column")
+
+    filters = []
+    if start_date:
+        filters.append(('date', '>=', start_date))
+    if end_date:
+        filters.append(('date', '<=', end_date))
+
+    if not identifiers:
+        # Load the entire database
+        return process_partition(storage_provider, base_path, identifier_column, columns=columns, filters=filters)
+    else:
+        identifiers = map_identifier(identifiers, input_identifier='ticker', output_identifier=identifier_column)
+
+    # Use numpy.isscalar to check if identifiers is a single value
+    if np.isscalar(identifiers):
+        # Single identifier
+        identifier = identifiers
+        return process_partition(storage_provider, base_path, identifier_column, identifier=identifier, columns=columns, filters=filters)
+    
+    # Multiple identifiers
+    with ThreadPoolExecutor() as executor:
+        futures = []
+        for identifier in identifiers:
+            futures.append(executor.submit(process_partition, storage_provider, base_path, identifier_column, identifier, columns, filters))
+        
+        results = []
+        for future in as_completed(futures):
+            results.append(future.result())
+    
+    df = pd.concat(results, ignore_index=True)
+    return df
+
+    # if 'filing_date' in df.columns:
+    #     df['date'] = pd.to_datetime(df['filing_date'])
+    # elif 'date' not in df.columns:
+    #     raise ValueError("Neither 'filing_date' nor 'date' column found in the data")
+    
+    # df.set_index([identifier_column, 'date'], inplace=True)
+    # df.sort_index(inplace=True)
+    
+    return df
+
+import pandas as pd
+import numpy as np
+
+def map_identifier(input_values, input_identifier='ticker', output_identifier='cik', trailing_zero_removal=True):
+    """
+    Map from one identifier to another using the df_codes DataFrame.
+    Handles string, single-item list, and multiple-item lists as input.
+    Uses vectorization for improved performance.
+    
+    :param input_values: The input identifier value(s) to map from (string, list, or array-like)
+    :param input_identifier: The type of input identifier (default: 'ticker')
+    :param output_identifier: The type of output identifier (default: 'cik')
+    :param trailing_zero_removal: If True, convert CIK to int (default: True)
+    :return: Mapped identifier value(s), same type as input
+    """
+    df_codes = pd.read_parquet("data/codes.parq")
+
+    # Check if the identifier columns exist in df_codes
+    if input_identifier not in df_codes.columns or output_identifier not in df_codes.columns:
+        raise ValueError(f"Identifier '{input_identifier}' or '{output_identifier}' not found in the DataFrame")
+    
+    # Convert input to numpy array for consistent handling
+    if isinstance(input_values, str):
+        input_array = np.array([input_values])
+        input_is_string = True
+    else:
+        input_array = np.array(input_values)
+        input_is_string = False
+    
+    # Ensure all input values are strings
+    input_array = np.char.strip(input_array.astype(str))
+    
+    # Create a dictionary for fast lookup
+    mapping_dict = dict(zip(df_codes[input_identifier], df_codes[output_identifier]))
+    
+    # Vectorized mapping
+    result = np.vectorize(mapping_dict.get)(input_array)
+    
+    # Convert to int if trailing_zero_removal is True and output is CIK
+    if trailing_zero_removal and output_identifier.lower() == 'cik':
+        result = np.where(result == None, 0, result)  # Replace None with 0
+        result = result.astype(np.int32)  # Convert to int64
+    
+    # Return result in the same format as input
+    if input_is_string:
+        return result[0]
+    elif len(input_array) == 1:
+        return result[0]
+    else:
+        return result.tolist()
+
+def load_frame_s3(endpoint, tickers=None, columns=None, start_date=None, end_date=None):
+
+    endpoint_config = {
+        "sec/10k": {
+            "storage_provider": {
+                "digitalocean": "sovai/sovai-filings/tenk",
+                "wasabi": "sovai-filings/tenk"
+            },
+            "identifier_column": "cik",
+            "partitioned": True,
+    
+        },
+        "future/enpdoint": {
+            "storage_provider": {
+                "digitalocean": "sovai/sovai-filings/tenk",
+                "wasabi": "sovai-filings/tenk"
+            },
+            "identifier_column": "cik",
+            "partitioned": True,
+        }
+    }
+    
+    if endpoint not in endpoint_config:
+        raise ValueError(f"Invalid endpoint: {endpoint}")
+    
+    config = endpoint_config[endpoint]
+
+    
+    if endpoint == "sec/10k":
+        df_frame = client_side_s3_frame(
+            config,
+            tickers,
+            columns,
+            start_date,
+            end_date
+        )
+    else:
+        # Implement logic for other endpoints if needed
+        raise NotImplementedError(f"Endpoint {endpoint} is not yet implemented")
+    
+    if HAS_CUSTOM_DATAFRAME:
+        return CustomDataFrame(df_frame)
+    else:
+        return df_frame  # Returns a regular pandas DataFrame if CustomDataFrame is not available
+