polarfrost 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

polarfrost/__init__.py

@@ -6,21 +6,78 @@ including the Mondrian algorithm, with support for both local (Polars)
 and distributed (PySpark) processing.
 """
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 # Import main functions
 try:
     from .mondrian import (
         mondrian_k_anonymity,
+        mondrian_k_anonymity_alt,
         mondrian_k_anonymity_polars,
-        mondrian_k_anonymity_spark
+        mondrian_k_anonymity_spark,
     )
+    from .clustering import clustering_k_anonymity
+
     __all__ = [
-        'mondrian_k_anonymity',
-        'mondrian_k_anonymity_polars',
-        'mondrian_k_anonymity_spark'
+        "mondrian_k_anonymity",
+        "mondrian_k_anonymity_alt",
+        "mondrian_k_anonymity_polars",
+        "mondrian_k_anonymity_spark",
+        "clustering_k_anonymity",
     ]
 except ImportError as e:
     import warnings
-    warnings.warn(f"Could not import mondrian: {e}")
+
+    warnings.warn(f"Could not import all dependencies: {e}")
+
+    # Define dummy functions for type checking
+    from typing import Any, List, Optional, Union
+    from typing_extensions import Literal
+    from pyspark.sql.types import StructType
+    from pyspark.sql import DataFrame as SparkDataFrame
+    from polars import DataFrame as PolarsDataFrame, LazyFrame
+    
+    def mondrian_k_anonymity(
+        df: Union[PolarsDataFrame, LazyFrame, SparkDataFrame],
+        quasi_identifiers: List[str],
+        sensitive_column: str,
+        k: int,
+        categorical: Optional[List[str]] = None,
+        schema: Optional[StructType] = None
+    ) -> Union[PolarsDataFrame, SparkDataFrame]:
+        """Dummy function for type checking when dependencies are missing."""
+        raise ImportError("Mondrian k-anonymity is not available due to missing dependencies")
+
+    def mondrian_k_anonymity_polars(
+        df: Union[PolarsDataFrame, LazyFrame],
+        quasi_identifiers: List[str],
+        sensitive_column: str,
+        k: int,
+        categorical: Optional[List[str]] = None
+    ) -> PolarsDataFrame:
+        """Dummy function for type checking when dependencies are missing."""
+        raise ImportError("Mondrian k-anonymity (Polars) is not available due to missing dependencies")
+
+    def mondrian_k_anonymity_spark(
+        df: SparkDataFrame,
+        quasi_identifiers: List[str],
+        sensitive_column: str,
+        k: int,
+        categorical: Optional[List[str]] = None,
+        schema: Optional[StructType] = None
+    ) -> SparkDataFrame:
+        """Dummy function for type checking when dependencies are missing."""
+        raise ImportError("Mondrian k-anonymity (PySpark) is not available due to missing dependencies")
+
+    def clustering_k_anonymity(
+        df: Union[PolarsDataFrame, LazyFrame],
+        quasi_identifiers: List[str],
+        sensitive_column: str,
+        k: int,
+        categorical: Optional[List[str]] = None,
+        method: str = "kmeans"
+    ) -> PolarsDataFrame:
+        """Dummy function for type checking when dependencies are missing."""
+        raise ImportError("Clustering k-anonymity is not available due to missing dependencies")
+
     __all__ = []

polarfrost/clustering.py

@@ -1,9 +1,19 @@
 """
 Clustering-based k-anonymity implementation using Polars.
 """
-from typing import List, Optional, Union
+
+from typing import List, Optional, Union, cast, TypeVar, Any
 import polars as pl
 
+# Define type variable for Polars expressions
+PolarsExpr = TypeVar('PolarsExpr', bound='pl.Expr')
+
+
+class ClusteringError(ValueError):
+    """Custom exception for clustering-related errors."""
+    pass
+
+
 def clustering_k_anonymity(
     df: Union[pl.DataFrame, pl.LazyFrame],
     quasi_identifiers: List[str],
@@ -14,16 +24,58 @@ def clustering_k_anonymity(
 ) -> pl.DataFrame:
     """
     Perform clustering-based k-anonymity using Polars.
-    
+
     Args:
         df: Input DataFrame or LazyFrame
         quasi_identifiers: List of column names to use for clustering
         sensitive_column: Column containing sensitive information
         k: Minimum group size for k-anonymity
-        categorical: List of categorical column names
-        method: Clustering method ('fcbg', 'rsc', or 'random')
-        
+        categorical: List of categorical column names (default: None)
+        method: Clustering method ('fcbg', 'rsc', or 'random') (default: 'fcbg')
+
     Returns:
         Anonymized DataFrame with generalized quasi-identifiers
+
+    Raises:
+        ValueError: If input validation fails
+        NotImplementedError: If the method is not implemented yet
     """
+    # Input validation
+    if not isinstance(df, (pl.DataFrame, pl.LazyFrame)):
+        raise ValueError("Input must be a Polars DataFrame or LazyFrame")
+    
+    # Validate quasi_identifiers
+    if not isinstance(quasi_identifiers, list) or not quasi_identifiers:
+        raise ValueError("quasi_identifiers must be a non-empty list")
+    
+    # Convert to LazyFrame if not already
+    is_lazy = isinstance(df, pl.LazyFrame)
+    if not is_lazy:
+        df = df.lazy()
+    
+    # Check for empty DataFrame
+    df_len = df.select(pl.len()).collect().item(0, 0) if not is_lazy else df.select(pl.len()).collect().item(0, 0)
+    if df_len == 0:
+        raise ValueError("Input DataFrame cannot be empty")
+
+    # Validate k is a positive integer
+    if not isinstance(k, (int, str)) or (isinstance(k, str) and not str(k).isdigit()) or int(k) < 1:
+        raise ValueError("k must be a positive integer")
+    k = int(k)  # Convert to int if it's a string of digits
+
+    # Get all columns that will be used
+    all_columns = set(quasi_identifiers + [sensitive_column])
+    if categorical is not None:
+        all_columns.update(categorical)
+    
+    # Check if columns exist
+    missing_columns = [col for col in all_columns if col not in df.columns]
+    if missing_columns:
+        raise ValueError(f"Columns not found in DataFrame: {missing_columns}")
+    
+    # Validate method
+    if method not in ["fcbg", "rsc", "random"]:
+        raise ValueError(f"Unsupported clustering method: {method}")
+    
+    # For now, just raise NotImplementedError
     raise NotImplementedError("Clustering k-anonymity will be implemented soon")