ygg 0.1.31__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,7 +1,10 @@
+"""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
@@ -452,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)
 
@@ -492,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
 
@@ -741,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,
@@ -757,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)
 
@@ -770,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
 
@@ -846,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)
@@ -1100,10 +1166,39 @@ def record_batch_reader_to_record_batch(
 def arrow_dataset_to_table(
     data: pds.Dataset,
     options: Optional[CastOptions] = None,
-) -> pa.Field:
+) -> 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
 # ---------------------------------------------------------------------------
@@ -1154,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")
@@ -1166,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()
 
@@ -1181,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)

yggdrasil/types/cast/cast_options.py

@@ -1,3 +1,5 @@
+"""Casting options for Arrow- and engine-aware conversions."""
+
 import dataclasses
 from typing import Optional, Union, List, Any
 
@@ -69,6 +71,22 @@ class CastOptions:
         target_field: pa.Field | pa.Schema | pa.DataType | None = None,
         **kwargs
     ):
+        """Build a CastOptions instance with optional source/target fields.
+
+        Args:
+            safe: Enable safe casting if True.
+            add_missing_columns: Add missing columns if True.
+            strict_match_names: Require exact field name matches if True.
+            allow_add_columns: Allow extra columns if True.
+            eager: Enable eager casting behavior if True.
+            datetime_patterns: Optional datetime parsing patterns.
+            source_field: Optional source Arrow field/schema/type.
+            target_field: Optional target Arrow field/schema/type.
+            **kwargs: Additional CastOptions fields.
+
+        Returns:
+            CastOptions instance.
+        """
         built = CastOptions(
             safe=safe,
             add_missing_columns=add_missing_columns,
@@ -169,6 +187,14 @@ class CastOptions:
         return result
 
     def check_source(self, obj: Any):
+        """Set the source field if not already configured.
+
+        Args:
+            obj: Source object to infer from.
+
+        Returns:
+            Self.
+        """
         if self.source_field is not None or obj is None:
             return self
 
@@ -177,6 +203,14 @@ class CastOptions:
         return self
 
     def need_arrow_type_cast(self, source_obj: Any):
+        """Return True when Arrow type casting is required.
+
+        Args:
+            source_obj: Source object to compare types against.
+
+        Returns:
+            True if Arrow type cast needed.
+        """
         if self.target_field is None:
             return False
 
@@ -185,6 +219,14 @@ class CastOptions:
         return self.source_field.type != self.target_field.type
 
     def need_polars_type_cast(self, source_obj: Any):
+        """Return True when Polars dtype casting is required.
+
+        Args:
+            source_obj: Source object to compare types against.
+
+        Returns:
+            True if Polars type cast needed.
+        """
         if self.target_polars_field is None:
             return False
 
@@ -193,6 +235,14 @@ class CastOptions:
         return self.source_polars_field.dtype != self.target_polars_field.dtype
 
     def need_spark_type_cast(self, source_obj: Any):
+        """Return True when Spark datatype casting is required.
+
+        Args:
+            source_obj: Source object to compare types against.
+
+        Returns:
+            True if Spark type cast needed.
+        """
         if self.target_spark_field is None:
             return False
 
@@ -201,6 +251,14 @@ class CastOptions:
         return self.source_spark_field.dataType != self.target_spark_field.dataType
 
     def need_nullability_check(self, source_obj: Any):
+        """Return True when nullability checks are required.
+
+        Args:
+            source_obj: Source object to compare nullability against.
+
+        Returns:
+            True if nullability check needed.
+        """
         if self.target_field is None:
             return False
 
@@ -213,6 +271,15 @@ class CastOptions:
         arrow_field: pa.Field,
         index: int
     ):
+        """Return a child Arrow field by index for nested types.
+
+        Args:
+            arrow_field: Parent Arrow field.
+            index: Child index.
+
+        Returns:
+            Child Arrow field.
+        """
         source_type: Union[
             pa.DataType, pa.ListType, pa.StructType, pa.MapType
         ] = arrow_field.type
@@ -235,6 +302,11 @@ class CastOptions:
 
     @property
     def source_field(self):
+        """Return the configured source Arrow field.
+
+        Returns:
+            Source Arrow field.
+        """
         return self.source_arrow_field
 
     @source_field.setter
@@ -248,10 +320,23 @@ class CastOptions:
         object.__setattr__(self, "source_arrow_field", value)
 
     def source_child_arrow_field(self, index: int):
+        """Return a child source Arrow field by index.
+
+        Args:
+            index: Child index.
+
+        Returns:
+            Child Arrow field.
+        """
         return self._child_arrow_field(self.source_arrow_field, index=index)
 
     @property
     def source_polars_field(self):
+        """Return or compute the cached Polars field for the source.
+
+        Returns:
+            Polars field or None.
+        """
         if self.source_arrow_field is not None and self._source_polars_field is None:
             from ...types.cast.polars_cast import arrow_field_to_polars_field
 
@@ -260,6 +345,11 @@ class CastOptions:
 
     @property
     def source_spark_field(self):
+        """Return or compute the cached Spark field for the source.
+
+        Returns:
+            Spark field or None.
+        """
         if self.source_arrow_field is not None and self._source_spark_field is None:
             from ...types.cast.spark_cast import arrow_field_to_spark_field
 
@@ -275,6 +365,11 @@ class CastOptions:
 
     @property
     def target_field_name(self):
+        """Return the effective target field name.
+
+        Returns:
+            Target field name or None.
+        """
         if self.target_field is None:
             if self.source_field is not None:
                 return self.source_field.name
@@ -295,10 +390,23 @@ class CastOptions:
         object.__setattr__(self, "target_arrow_field", value)
 
     def target_child_arrow_field(self, index: int):
+        """Return a child target Arrow field by index.
+
+        Args:
+            index: Child index.
+
+        Returns:
+            Child Arrow field.
+        """
         return self._child_arrow_field(self.target_arrow_field, index=index)
 
     @property
     def target_polars_field(self):
+        """Return or compute the cached Polars field for the target.
+
+        Returns:
+            Polars field or None.
+        """
         if self.target_arrow_field is not None and self._target_polars_field is None:
             from ...types.cast.polars_cast import arrow_field_to_polars_field
 
@@ -307,6 +415,11 @@ class CastOptions:
 
     @property
     def target_spark_field(self):
+        """Return or compute the cached Spark field for the target.
+
+        Returns:
+            Spark field or None.
+        """
         if self.target_arrow_field is not None and self._target_spark_field is None:
             from ...types.cast.spark_cast import arrow_field_to_spark_field
 
@@ -329,6 +442,11 @@ class CastOptions:
 
     @property
     def target_spark_schema(self) -> Optional["pyspark.sql.types.StructType"]:
+        """Return a Spark schema view of the target Arrow schema.
+
+        Returns:
+            Spark StructType schema or None.
+        """
         arrow_schema = self.target_arrow_schema
 
         if arrow_schema is not None:
@@ -338,4 +456,4 @@ class CastOptions:
         return arrow_schema
 
 
-DEFAULT_INSTANCE = CastOptions()
+DEFAULT_INSTANCE = CastOptions()