Export cache

[lmmx]

You might want to export a cache (remove symlinks). Could this be done more easily by reversing the symlinks?

[lmmx]

Maybe better to use the filename callback and then exporting is trivial, rather than renaming everything from a concatenated args-based filename...

from enum import Enum
from typing import Callable
import urllib.parse

class SymlinkStyle(Enum):
    """Predefined symlink naming styles for different use cases."""
    
    DEFAULT = "default"           # output.parquet
    FUNCTION_ARGS = "func_args"   # load_data_sales_2024_north.parquet  
    ARGS_ONLY = "args_only"       # sales_2024_north.parquet
    TIMESTAMPED = "timestamped"   # load_data_20241201_143052.parquet
    NUMBERED = "numbered"         # load_data_001.parquet
    DATASET = "dataset"           # data.parquet (simple, for dataset-like functions)


def _create_symlink_name_from_style(
    style: SymlinkStyle,
    func: Callable,
    args: tuple,
    kwargs: dict,
    result,
    cache_key: str
) -> str:
    """Generate symlink filename based on the chosen style."""
    
    if style == SymlinkStyle.DEFAULT:
        return "output.parquet"
    
    elif style == SymlinkStyle.FUNCTION_ARGS:
        # function_name_arg1_arg2_kwarg1_kwarg2.parquet
        parts = [func.__name__]
        
        # Add positional args (cleaned)
        for arg in args:
            clean_arg = str(arg).replace('/', '_').replace(' ', '_')[:20]
            parts.append(clean_arg)
        
        # Add keyword args
        for key, value in kwargs.items():
            clean_value = str(value).replace('/', '_').replace(' ', '_')[:20]
            parts.append(f"{key}_{clean_value}")
        
        return "_".join(parts) + ".parquet"
    
    elif style == SymlinkStyle.ARGS_ONLY:
        # Just the arguments: sales_2024_north.parquet
        parts = []
        
        for arg in args:
            clean_arg = str(arg).replace('/', '_').replace(' ', '_')[:20]
            parts.append(clean_arg)
        
        for key, value in kwargs.items():
            clean_value = str(value).replace('/', '_').replace(' ', '_')[:20]
            parts.append(f"{key}_{clean_value}")
        
        if parts:
            return "_".join(parts) + ".parquet"
        else:
            return "data.parquet"
    
    elif style == SymlinkStyle.TIMESTAMPED:
        # function_name_20241201_143052.parquet
        import time
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        return f"{func.__name__}_{timestamp}.parquet"
    
    elif style == SymlinkStyle.NUMBERED:
        # function_name_001.parquet (would need counter logic)
        # For now, just use cache key suffix
        short_key = cache_key[:6]
        return f"{func.__name__}_{short_key}.parquet"
    
    elif style == SymlinkStyle.DATASET:
        # Simple data.parquet - good for dataset loading functions
        return "data.parquet"
    
    else:
        return "output.parquet"


# Update the PolarsCache class to accept SymlinkStyle
class PolarsCache:
    def __init__(
        self,
        cache_dir: str | None = None,
        use_tmp: bool = False,
        hidden: bool = True,
        size_limit: int | str = "1GB",
        symlinks_dir: str = "functions",
        nested: bool = True,
        trim_arg: int = 50,
        symlink_name: str | FilenameCallback | SymlinkStyle | None = None,
    ):
        # ... existing init code ...
        
        # Convert SymlinkStyle enum to callable
        if isinstance(symlink_name, SymlinkStyle):
            self.symlink_name = lambda f, a, k, r, c: _create_symlink_name_from_style(
                symlink_name, f, a, k, r, c
            )
        else:
            self.symlink_name = symlink_name


# Update the convenience cache function too
def cache(
    cache_dir: str | None = None,
    use_tmp: bool = False,
    hidden: bool = True,
    size_limit: int | str = "1GB",
    symlinks_dir: str = "functions",
    nested: bool = True,
    trim_arg: int = 50,
    symlink_name: str | SymlinkStyle | None = None,
):
    # ... rest of function with SymlinkStyle support


# Usage examples:

# Example 1: Clean function+args naming
@cache(symlink_name=SymlinkStyle.FUNCTION_ARGS)
def load_sales_data(year: int, region: str) -> pl.DataFrame:
    return pl.DataFrame({"sales": [100, 200], "year": [year, year]})

# Creates: functions/module/load_sales_data/arg0=2024_arg1=north/load_sales_data_2024_north.parquet
load_sales_data(2024, "north")

# Example 2: Just arguments (good for dataset functions) 
@cache(symlink_name=SymlinkStyle.ARGS_ONLY)
def get_dataset(name: str, version: str = "latest") -> pl.DataFrame:
    return pl.DataFrame({"data": [1, 2, 3]})

# Creates: functions/module/get_dataset/arg0=mnist_version=v2/mnist_version_v2.parquet
get_dataset("mnist", version="v2")

# Example 3: Simple dataset naming
@cache(symlink_name=SymlinkStyle.DATASET)
def load_training_data() -> pl.DataFrame:
    return pl.DataFrame({"features": [1, 2, 3]})

# Creates: functions/module/load_training_data/no_args/data.parquet
load_training_data()

# Example 4: Timestamped (good for experiments)
@cache(symlink_name=SymlinkStyle.TIMESTAMPED)  
def run_experiment(learning_rate: float) -> pl.DataFrame:
    return pl.DataFrame({"loss": [0.5, 0.3, 0.1]})

# Creates: functions/module/run_experiment/arg0=0.01/run_experiment_20241201_143052.parquet
run_experiment(0.01)


# Now export becomes trivial:
def export_cache(self, export_dir: str | Path) -> int:
    """Export cache - just copy files, names are already good!"""
    export_path = Path(export_dir)
    export_path.mkdir(parents=True, exist_ok=True)
    
    count = 0
    for symlink_file in self.readable_dir.rglob("*.parquet"):
        if symlink_file.is_symlink():
            try:
                # Just use the symlink filename as-is!
                export_file = export_path / symlink_file.name
                
                # Handle collisions by adding parent dir
                if export_file.exists():
                    parent_name = symlink_file.parent.name
                    name_without_ext = symlink_file.stem
                    export_file = export_path / f"{parent_name}_{name_without_ext}.parquet"
                
                import shutil
                shutil.copy2(symlink_file, export_file)
                count += 1
            except (OSError, FileNotFoundError):
                continue
    
    return count


# The exported files would have clean names like:
# exported/
# ├── load_sales_data_2024_north.parquet
# ├── mnist_version_v2.parquet  
# ├── data.parquet
# └── run_experiment_20241201_143052.parquet