ygg 0.1.30__py3-none-any.whl → 0.1.32__py3-none-any.whl

yggdrasil/requests/msal.py

@@ -1,3 +1,5 @@
+"""MSAL-backed authentication helpers for requests sessions."""
+
 # auth_session.py
 import os
 import time
@@ -27,6 +29,15 @@ __all__ = [
 
 @dataclass
 class MSALAuth:
+    """Configuration and token cache for MSAL client credential flows.
+
+    Args:
+        tenant_id: Azure tenant ID.
+        client_id: Azure application client ID.
+        client_secret: Azure application client secret.
+        authority: Optional authority URL override.
+        scopes: List of scopes to request.
+    """
     tenant_id: Optional[str] = None
     client_id: Optional[str] = None
     client_secret: Optional[str] = None
@@ -38,12 +49,34 @@ class MSALAuth:
     _access_token: Optional[str] = None
 
     def __setitem__(self, key, value):
+        """Set an attribute via mapping-style assignment.
+
+        Args:
+            key: Attribute name to set.
+            value: Value to assign.
+
+        Returns:
+            None.
+        """
         self.__setattr__(key, value)
 
     def __getitem__(self, item):
+        """Return attribute values via mapping-style access.
+
+        Args:
+            item: Attribute name to fetch.
+
+        Returns:
+            The attribute value.
+        """
         return getattr(self, item)
 
     def __post_init__(self):
+        """Populate defaults from environment variables and validate.
+
+        Returns:
+            None.
+        """
         self.tenant_id = self.tenant_id or os.environ.get("AZURE_TENANT_ID")
         self.client_id = self.client_id or os.environ.get("AZURE_CLIENT_ID")
         self.client_secret = self.client_secret or os.environ.get("AZURE_CLIENT_SECRET")
@@ -60,7 +93,11 @@ class MSALAuth:
         self._validate_config()
 
     def _validate_config(self):
-        """Validate that all required configuration is present."""
+        """Validate that all required configuration is present.
+
+        Returns:
+            None.
+        """
         missing = []
 
         if not self.client_id:
@@ -81,6 +118,15 @@ class MSALAuth:
         env: Mapping = None,
         prefix: Optional[str] = None
     ) -> "MSALAuth":
+        """Return an MSALAuth built from environment variables if available.
+
+        Args:
+            env: Mapping to read variables from; defaults to os.environ.
+            prefix: Optional prefix for variable names.
+
+        Returns:
+            A configured MSALAuth instance or None.
+        """
         if not env:
             env = os.environ
         prefix = prefix or "AZURE_"
@@ -105,6 +151,14 @@ class MSALAuth:
         return None
 
     def export_to(self, to: dict = os.environ):
+        """Export the auth configuration to the provided mapping.
+
+        Args:
+            to: Mapping to populate with auth configuration values.
+
+        Returns:
+            None.
+        """
         for key, value in (
             ("AZURE_CLIENT_ID", self.client_id),
             ("AZURE_CLIENT_SECRET", self.client_secret),
@@ -116,6 +170,11 @@ class MSALAuth:
 
     @property
     def auth_app(self) -> ConfidentialClientApplication:
+        """Return or initialize the MSAL confidential client.
+
+        Returns:
+            MSAL confidential client instance.
+        """
         if not self._auth_app:
             self._auth_app = ConfidentialClientApplication(
                 client_id=self.client_id,
@@ -127,19 +186,42 @@ class MSALAuth:
 
     @property
     def expires_in(self) -> float:
+        """Return the number of seconds since the token expiry timestamp.
+
+        Returns:
+            Seconds elapsed since expiry (negative if not expired).
+        """
         return time.time() - self.expires_at
 
     @property
     def expires_at(self) -> float:
+        """Ensure the token is fresh and return the expiry timestamp.
+
+        Returns:
+            Token expiration time as a Unix timestamp.
+        """
         self.refresh()
 
         return self._expires_at
 
     @property
     def expired(self) -> bool:
+        """Return True when the token is missing or past its expiry time.
+
+        Returns:
+            True if expired or missing; False otherwise.
+        """
         return not self._expires_at or time.time() >= self._expires_at
 
     def refresh(self, force: bool | None = None):
+        """Acquire or refresh the token if needed.
+
+        Args:
+            force: Force refresh even if not expired.
+
+        Returns:
+            The updated MSALAuth instance.
+        """
         if self.expired or force:
             app = self.auth_app
             result = app.acquire_token_for_client(scopes=self.scopes)
@@ -157,16 +239,32 @@ class MSALAuth:
 
     @property
     def access_token(self) -> str:
-        """Return access token."""
+        """Return access token.
+
+        Returns:
+            Access token string.
+        """
         self.refresh()
         return self._access_token
 
     @property
     def authorization(self) -> str:
-        """Return authorization token."""
+        """Return authorization token.
+
+        Returns:
+            Authorization header value.
+        """
         return f"Bearer {self.access_token}"
 
     def requests_session(self, **kwargs):
+        """Build a requests session that injects the MSAL authorization header.
+
+        Args:
+            **kwargs: Passed through to MSALSession.
+
+        Returns:
+            Configured MSALSession.
+        """
         return MSALSession(
             msal_auth=self,
             **kwargs
@@ -174,6 +272,11 @@ class MSALAuth:
 
 
 class MSALSession(YGGSession):
+    """YGGSession subclass that injects MSAL authorization headers.
+
+    Args:
+        YGGSession: Base retry-capable session.
+    """
     msal_auth: MSALAuth | None = None
 
     def __init__(
@@ -182,11 +285,29 @@ class MSALSession(YGGSession):
         *args,
         **kwargs: dict
     ):
+        """Initialize the session with optional MSAL auth configuration.
+
+        Args:
+            msal_auth: MSALAuth configuration for token injection.
+            *args: Positional args for YGGSession.
+            **kwargs: Keyword args for YGGSession.
+
+        Returns:
+            None.
+        """
         super().__init__(*args, **kwargs)
         self.msal_auth = msal_auth
 
 
     def prepare_request(self, request):
+        """Prepare the request with an Authorization header when needed.
+
+        Args:
+            request: requests.PreparedRequest to mutate.
+
+        Returns:
+            Prepared request.
+        """
         # called before sending; ensure header exists
         if self.msal_auth is not None:
             request.headers["Authorization"] = request.headers.get("Authorization", self.msal_auth.authorization)

yggdrasil/requests/session.py

@@ -1,3 +1,5 @@
+"""HTTP session helpers with retry-enabled defaults."""
+
 from typing import Optional, Dict
 
 from requests import Session
@@ -10,6 +12,11 @@ __all__ = [
 
 
 class YGGSession(Session):
+    """Requests session with preconfigured retry adapter support.
+
+    Args:
+        Session: Base requests session type.
+    """
     def __init__(
         self,
         num_retry: int = 4,
@@ -17,6 +24,17 @@ class YGGSession(Session):
         *args,
         **kwargs
     ):
+        """Initialize the session with retries and optional default headers.
+
+        Args:
+            num_retry: Number of retries for connection and read errors.
+            headers: Optional default headers to merge into the session.
+            *args: Additional positional arguments passed to Session.
+            **kwargs: Additional keyword arguments passed to Session.
+
+        Returns:
+            None.
+        """
         super(YGGSession, self).__init__()
 
         retry = Retry(

yggdrasil/types/__init__.py

@@ -1,3 +1,5 @@
+"""Type utilities for Arrow inference and casting."""
+
 from .python_arrow import *
 from .python_defaults import *
 from .cast import *

yggdrasil/types/cast/__init__.py

@@ -1,3 +1,5 @@
+"""Casting utilities and converters across Arrow and engine types."""
+
 from .registry import *
 from .arrow_cast import *
 from .polars_cast import *
@@ -6,4 +8,3 @@ from .spark_cast import *
 from .spark_polars_cast import *
 from .spark_pandas_cast import *
 from .polars_pandas_cast import *
-

yggdrasil/types/cast/arrow_cast.py

@@ -1,11 +1,15 @@
+"""Arrow casting helpers for arrays, tables, and schemas."""
+
 import dataclasses
 import enum
 import logging
 from dataclasses import is_dataclass
+from functools import partial
 from typing import Optional, Union, List, Tuple, Any
 
 import pyarrow as pa
 import pyarrow.compute as pc
+import pyarrow.dataset as pds
 
 from .cast_options import CastOptions
 from .registry import register_converter
@@ -451,6 +455,15 @@ def any_to_arrow_scalar(
     scalar: Any,
     options: Optional[CastOptions] = None,
 ) -> pa.Scalar:
+    """Convert a Python value to an Arrow scalar.
+
+    Args:
+        scalar: Input value.
+        options: Optional cast options.
+
+    Returns:
+        Arrow scalar.
+    """
     if isinstance(scalar, pa.Scalar):
         return cast_arrow_scalar(scalar, options)
 
@@ -491,6 +504,15 @@ def cast_arrow_scalar(
     scalar: pa.Scalar,
     options: Optional[CastOptions] = None,
 ) -> pa.Scalar:
+    """Cast an Arrow scalar to the target Arrow field.
+
+    Args:
+        scalar: Arrow scalar to cast.
+        options: Optional cast options.
+
+    Returns:
+        Casted Arrow scalar.
+    """
     options = CastOptions.check_arg(options)
     target_field = options.target_field
 
@@ -740,6 +762,28 @@ def cast_arrow_tabular(
     return data.__class__.from_arrays(all_arrays, schema=target_arrow_schema)
 
 
+@register_converter(pds.Dataset, pds.Dataset)
+def cast_arrow_dataset(data: pds.Dataset, options: Optional[CastOptions] = None) -> pds.Dataset:
+    """Cast a dataset to the target schema in options.
+
+    Args:
+        data: Arrow dataset to cast.
+        options: Optional cast options.
+
+    Returns:
+        Casted dataset.
+    """
+    if options is None:
+        return data
+
+    caster = partial(cast_arrow_tabular, options=options)
+
+    return pds.dataset(map(caster, data.to_batches(
+        batch_size=256 * 1024,
+        memory_pool=options.arrow_memory_pool
+    )))
+
+
 @register_converter(pa.RecordBatchReader, pa.RecordBatchReader)
 def cast_arrow_record_batch_reader(
     data: pa.RecordBatchReader,
@@ -756,6 +800,11 @@ def cast_arrow_record_batch_reader(
         return data
 
     def casted_batches():
+        """Yield casted batches from a RecordBatchReader.
+
+        Yields:
+            Casted RecordBatch instances.
+        """
         for batch in data:
             yield cast_arrow_tabular(batch, options)
 
@@ -769,6 +818,15 @@ def any_to_arrow_array(
     obj: Any,
     options: Optional[CastOptions] = None,
 ) -> pa.Array:
+    """Convert array-like input into an Arrow array.
+
+    Args:
+        obj: Array-like input.
+        options: Optional cast options.
+
+    Returns:
+        Arrow array.
+    """
     options = CastOptions.check_arg(options)
     arrow_array = None
 
@@ -845,6 +903,15 @@ def pylist_to_record_batch(
     data: list,
     options: Optional[CastOptions] = None,
 ) -> pa.RecordBatch:
+    """Convert a list of rows into a RecordBatch.
+
+    Args:
+        data: List of row objects.
+        options: Optional cast options.
+
+    Returns:
+        Arrow RecordBatch.
+    """
     options = CastOptions.check_arg(options)
 
     array: Union[pa.Array, pa.StructArray] = any_to_arrow_array(data, options)
@@ -1095,6 +1162,43 @@ def record_batch_reader_to_record_batch(
     return table_to_record_batch(table, options)
 
 
+@register_converter(pds.Dataset, pa.Table)
+def arrow_dataset_to_table(
+    data: pds.Dataset,
+    options: Optional[CastOptions] = None,
+) -> pa.Table:
+    """Convert a dataset to a Table and apply casting.
+
+    Args:
+        data: Arrow dataset.
+        options: Optional cast options.
+
+    Returns:
+        Arrow table.
+    """
+    table = data.to_table()
+    return cast_arrow_tabular(table, options)
+
+
+@register_converter(pa.Table, pds.Dataset)
+@register_converter(pa.RecordBatch, pds.Dataset)
+def arrow_tabular_to_dataset(
+    data: Union[pa.Table, pa.RecordBatch],
+    options: Optional[CastOptions] = None,
+) -> pa.Field:
+    """Convert Arrow tabular data to a dataset after casting.
+
+    Args:
+        data: Table or RecordBatch.
+        options: Optional cast options.
+
+    Returns:
+        Arrow dataset.
+    """
+    data = cast_arrow_tabular(data, options)
+    return pds.dataset([data])
+
+
 # ---------------------------------------------------------------------------
 # Field / Schema converters
 # ---------------------------------------------------------------------------
@@ -1145,6 +1249,15 @@ def arrow_schema_to_field(
     data: pa.Schema,
     options: Optional[CastOptions] = None,
 ) -> pa.Field:
+    """Wrap an Arrow schema as a struct field.
+
+    Args:
+        data: Arrow schema.
+        options: Optional cast options.
+
+    Returns:
+        Arrow field.
+    """
     dtype = pa.struct(list(data))
     md = dict(data.metadata or {})
     name = md.setdefault(b"name", b"root")
@@ -1157,6 +1270,15 @@ def arrow_field_to_schema(
     data: pa.Field,
     options: Optional[CastOptions] = None,
 ) -> pa.Schema:
+    """Return a schema view of an Arrow field.
+
+    Args:
+        data: Arrow field.
+        options: Optional cast options.
+
+    Returns:
+        Arrow schema.
+    """
     md = dict(data.metadata or {})
     md[b"name"] = data.name.encode()
 
@@ -1172,4 +1294,13 @@ def arrow_tabular_to_field(
     data: Union[pa.Table, pa.RecordBatch, pa.RecordBatchReader],
     options: Optional[CastOptions] = None,
 ) -> pa.Field:
+    """Return a field representing the schema of tabular data.
+
+    Args:
+        data: Arrow table/batch/reader.
+        options: Optional cast options.
+
+    Returns:
+        Arrow field.
+    """
     return arrow_schema_to_field(data.schema, options)