arize 8.0.0a21__py3-none-any.whl → 8.0.0a23__py3-none-any.whl

arize/models/bounded_executor.py

@@ -1,23 +1,35 @@
+"""Bounded thread pool executor with queue size limits."""
+
+from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 from threading import BoundedSemaphore
 
 
 class BoundedExecutor:
-    """
-    BoundedExecutor behaves as a ThreadPoolExecutor which will block on
-    calls to submit() once the limit given as "bound" work items are queued for
-    execution.
+    """BoundedExecutor behaves as a ThreadPoolExecutor which will block on calls to submit().
+
+    Blocks once the limit given as "bound" work items are queued for execution.
+
     :param bound: Integer - the maximum number of items in the work queue
     :param max_workers: Integer - the size of the thread pool
     """
 
-    def __init__(self, bound, max_workers):
+    def __init__(self, bound: int, max_workers: int) -> None:
+        """Initialize the bounded executor.
+
+        Args:
+            bound: Maximum number of items in the work queue.
+            max_workers: Size of the thread pool.
+        """
         self.executor = ThreadPoolExecutor(max_workers=max_workers)
         self.semaphore = BoundedSemaphore(bound + max_workers)
 
     """See concurrent.futures.Executor#submit"""
 
-    def submit(self, fn, *args, **kwargs):
+    def submit(
+        self, fn: Callable[..., object], *args: object, **kwargs: object
+    ) -> object:
+        """Submit a callable to be executed with bounded concurrency."""
         self.semaphore.acquire()
         try:
             future = self.executor.submit(fn, *args, **kwargs)
@@ -30,5 +42,6 @@ class BoundedExecutor:
 
     """See concurrent.futures.Executor#shutdown"""
 
-    def shutdown(self, wait=True):
+    def shutdown(self, wait: bool = True) -> None:
+        """Shutdown the executor, optionally waiting for pending tasks to complete."""
         self.executor.shutdown(wait)

arize/models/casting.py

@@ -1,8 +1,10 @@
+"""Type casting utilities for ML model data conversion."""
+
 # type: ignore[pb2]
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING, List, Tuple, Union
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -14,14 +16,24 @@ if TYPE_CHECKING:
 
 
 class CastingError(Exception):
+    """Raised when type casting fails for a value."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, error_msg: str, typed_value: TypedValue) -> None:
+        """Initialize the exception with type casting context.
+
+        Args:
+            error_msg: Description of the casting failure.
+            typed_value: The TypedValue that failed to cast.
+        """
         self.error_msg = error_msg
         self.typed_value = typed_value
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"Failed to cast value {self.typed_value.value} of type {type(self.typed_value.value)} "
             f"to type {self.typed_value.type}. "
@@ -30,7 +42,10 @@ class CastingError(Exception):
 
 
 class ColumnCastingError(Exception):
+    """Raised when type casting fails for a column."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(
@@ -39,11 +54,19 @@ class ColumnCastingError(Exception):
         attempted_columns: str,
         attempted_type: TypedColumns,
     ) -> None:
+        """Initialize the exception with column casting context.
+
+        Args:
+            error_msg: Description of the casting failure.
+            attempted_columns: Columns that failed to cast.
+            attempted_type: The TypedColumns type that was attempted.
+        """
         self.error_msg = error_msg
         self.attempted_casting_columns = attempted_columns
         self.attempted_casting_type = attempted_type
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"Failed to cast to type {self.attempted_casting_type} "
             f"for columns: {log_a_list(self.attempted_casting_columns, 'and')}. "
@@ -52,36 +75,55 @@ class ColumnCastingError(Exception):
 
 
 class InvalidTypedColumnsError(Exception):
+    """Raised when typed columns are invalid or incorrectly specified."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, field_name: str, reason: str) -> None:
+        """Initialize the exception with typed columns validation context.
+
+        Args:
+            field_name: Name of the schema field with invalid typed columns.
+            reason: Description of why the typed columns are invalid.
+        """
         self.field_name = field_name
         self.reason = reason
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return f"The {self.field_name} TypedColumns object {self.reason}."
 
 
 class InvalidSchemaFieldTypeError(Exception):
+    """Raised when schema field has invalid or unexpected type."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, msg: str) -> None:
+        """Initialize the exception with schema field type error message.
+
+        Args:
+            msg: Error message describing the schema field type issue.
+        """
         self.msg = msg
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return self.msg
 
 
 def cast_typed_columns(
     dataframe: pd.DataFrame,
     schema: Schema,
-) -> Tuple[pd.DataFrame, Schema]:
-    """
-    Cast feature and tag columns in the dataframe to the types specified in each TypedColumns config.
-    This optional feature provides a simple way for users to prevent
-    type drift within a column across many SDK uploads.
+) -> tuple[pd.DataFrame, Schema]:
+    """Cast feature and tag columns in the dataframe to the types specified in each TypedColumns config.
+
+    This optional feature provides a simple way for users to prevent type drift within
+    a column across many SDK uploads.
 
     Arguments:
     ---------
@@ -120,7 +162,7 @@ def cast_typed_columns(
         )
 
     # Make sure no other schema fields have this type.
-    if any({f for f in typed_column_fields if f not in allowed_fields}):
+    if any(f for f in typed_column_fields if f not in allowed_fields):
         raise InvalidSchemaFieldTypeError(
             "Only the feature_column_names and tag_column_names Schema fields can be of type "
             "TypedColumns. Fields with type TypedColumns:"
@@ -130,10 +172,7 @@ def cast_typed_columns(
     for field_name in typed_column_fields:
         f = getattr(schema, field_name)
         if f:
-            try:
-                _validate_typed_columns(field_name, f)
-            except InvalidTypedColumnsError:
-                raise
+            _validate_typed_columns(field_name, f)
             dataframe = _cast_columns(dataframe, f)
 
     # Now that the dataframe values have been cast to the specified types:
@@ -144,6 +183,14 @@ def cast_typed_columns(
 
 
 def cast_dictionary(d: dict) -> dict:
+    """Cast TypedValue entries in a dictionary to their appropriate Python types.
+
+    Args:
+        d: Dictionary that may contain TypedValue objects as values.
+
+    Returns:
+        Dictionary with TypedValue objects cast to their native Python types.
+    """
     cast_dict = {}
     for k, v in d.items():
         if isinstance(v, TypedValue):
@@ -154,9 +201,8 @@ def cast_dictionary(d: dict) -> dict:
 
 def _cast_value(
     typed_value: TypedValue,
-) -> Union[str, int, float, List[str], None]:
-    """
-    Casts a TypedValue to its provided type, preserving all null values as None or float('nan').
+) -> str | int | float | list[str] | None:
+    """Casts a TypedValue to its provided type, preserving all null values as None or float('nan').
 
     Arguments:
     ---------
@@ -179,22 +225,21 @@ def _cast_value(
 
     if typed_value.type == ArizeTypes.FLOAT:
         return _cast_to_float(typed_value)
-    elif typed_value.type == ArizeTypes.INT:
+    if typed_value.type == ArizeTypes.INT:
         return _cast_to_int(typed_value)
-    elif typed_value.type == ArizeTypes.STR:
+    if typed_value.type == ArizeTypes.STR:
         return _cast_to_str(typed_value)
-    else:
-        raise CastingError("Unknown casting type", typed_value)
+    raise CastingError("Unknown casting type", typed_value)
 
 
-def _cast_to_float(typed_value: TypedValue) -> Union[float, None]:
+def _cast_to_float(typed_value: TypedValue) -> float | None:
     try:
         return float(typed_value.value)
     except Exception as e:
         raise CastingError(str(e), typed_value) from e
 
 
-def _cast_to_int(typed_value: TypedValue) -> Union[int, None]:
+def _cast_to_int(typed_value: TypedValue) -> int | None:
     # a NaN float can't be cast to an int. Proactively return None instead.
     if isinstance(typed_value.value, float) and math.isnan(typed_value.value):
         return None
@@ -214,7 +259,7 @@ def _cast_to_int(typed_value: TypedValue) -> Union[int, None]:
         raise CastingError(str(e), typed_value) from e
 
 
-def _cast_to_str(typed_value: TypedValue) -> Union[str, None]:
+def _cast_to_str(typed_value: TypedValue) -> str | None:
     # a NaN float can't be cast to a string. Proactively return None instead.
     if isinstance(typed_value.value, float) and math.isnan(typed_value.value):
         return None
@@ -227,8 +272,7 @@ def _cast_to_str(typed_value: TypedValue) -> Union[str, None]:
 def _validate_typed_columns(
     field_name: str, typed_columns: TypedColumns
 ) -> None:
-    """
-    Validate a TypedColumns object.
+    """Validate a TypedColumns object.
 
     Arguments:
     ---------
@@ -256,8 +300,8 @@ def _validate_typed_columns(
 def _cast_columns(
     dataframe: pd.DataFrame, columns: TypedColumns
 ) -> pd.DataFrame:
-    """
-    Cast columns corresponding to a single TypedColumns object and a single Arize Schema field.
+    """Cast columns corresponding to a single TypedColumns object and a single Arize Schema field.
+
     (feature_column_names or tag_column_names)
 
     Arguments:
@@ -324,9 +368,10 @@ def _cast_columns(
 
 
 def _cast_df(
-    df: pd.DataFrame, cols: List[str], target_type_str: str
+    df: pd.DataFrame, cols: list[str], target_type_str: str
 ) -> pd.DataFrame:
-    """
+    """Cast columns in a dataframe to the specified type.
+
     Arguments:
     ---------
         df: pd.DataFrame
@@ -351,13 +396,14 @@ def _cast_df(
     df = df.replace(nan_mapping)
 
     # None or NaN-based values (including np.nan) are automatically converted to pandas pd.NA type
-    return df.astype({col: target_type_str for col in cols})
+    return df.astype(dict.fromkeys(cols, target_type_str))
 
 
 def _convert_schema_field_types(
     schema: Schema,
 ) -> Schema:
-    """
+    """Convert schema field types from TypedColumns to List[string] format.
+
     Arguments:
     ---------
         schema: Schema