nuthatch 0.1.0__py3-none-any.whl

nuthatch/cli.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""
+Command-line interface for nuthatch.
+
+This module provides a CLI for interacting with nuthatch caching functionality,
+including cache management, backend operations, and configuration.
+"""
+
+import importlib
+import click
+import shutil
+from .config import get_config
+from .backend import get_backend_by_name, registered_backends
+from .cache import Cache
+import pandas as pd
+
+
+@click.group()
+@click.version_option(version="0.1.0", prog_name="nuthatch")
+def cli():
+    """Nuthatch - Intelligent caching system for data science workflows.
+
+    This CLI provides tools for managing cache entries, inspecting backends,
+    and configuring the nuthatch system.
+    """
+
+    config = get_config(location='root', requested_parameters=['dynamic_config_path'], backend_name=None)
+    if 'dynamic_config_path' in config:
+        try:
+            importlib.import_module(config['dynamic_config_path'])
+        except Exception as e:
+            click.echo(f"WARNGIN: Failed to import {config['dynamic_config_path']} with '{e}'. You may be missing dynamic secret resolution.")
+
+
+
+@cli.command('import')
+@click.argument('cache_key')
+@click.option('--namespace', help='Namespace for the cache')
+@click.option('--backend', help='Backend to use', required=True)
+@click.option('--location', help='Location to search', default='root')
+def import_data(cache_key, namespace, backend, location):
+    """Import data from a glob pattern."""
+
+    # First instantiate the backend based on the passed backend
+    backend_name = backend
+    backend_class = get_backend_by_name(backend)
+    config = get_config(location=location, requested_parameters=backend_class.config_parameters, backend_name=backend_class.backend_name)
+    backend = backend_class(config, cache_key, namespace, None, {})
+
+    cache_keys = []
+    if hasattr(backend, 'fs') and backend.fs is not None:
+        paths = backend.fs.glob(backend.path)
+        for path in paths:
+            cache_keys.append(backend.get_cache_key(path))
+
+    if len(cache_keys) > 0:
+        click.confirm(f"Are you sure you want to import {len(paths)} cache entries?", abort=True)
+    else:
+        print("No caches found for import.")
+
+    for key in cache_keys:
+        print(f"Importing {key}.")
+
+        if backend_name == 'null':
+            config = get_config(location=location, requested_parameters=Cache.config_parameters)
+            cache = Cache(config, key, namespace, None, location, None, {})
+            if cache.is_null():
+                print(f"{key} already in cache as null!")
+            elif cache.exists():
+                print(f"Cache {key} already exists and is valid. Skipping entry. Delete this cache key if you would like to reimport it as null.")
+            else:
+                cache.set_null()
+                print(f"Set {key} successfully to null.")
+        else:
+            config = get_config(location=location, requested_parameters=Cache.config_parameters)
+            cache = Cache(config, key, namespace, None, location, backend_name, {})
+            if not cache.exists():
+                cache._commit_metadata()
+                print(f"Imported {key} successfully.")
+            else:
+                print(f"{key} already in cache!")
+
+def list_helper(cache_key, namespace, backend, location):
+    """List all cache entries."""
+    config = get_config(location=location, requested_parameters=Cache.config_parameters)
+    cache = Cache(config, None, namespace, None, location, backend, {})
+
+    if cache_key is None:
+        cache_key = '*'
+
+    caches = cache.list(cache_key)
+
+    return caches
+
+@cli.command('list')
+@click.argument('cache_key', required=False)
+@click.option('--namespace', help='Namespace for the cache')
+@click.option('--backend', help='Backend filter')
+@click.option('--location', help='Location to search', default='root')
+@click.option('--long', '-l', is_flag=True, help='List all information about the cache')
+def list_caches(cache_key, namespace, backend, location, long):
+
+    caches = list_helper(cache_key, namespace, backend, location)
+    pager = len(caches) > shutil.get_terminal_size()[0]
+
+    if not long:
+        caches = [cache['cache_key'] for cache in caches]
+        caches = '\n'.join(caches)
+    else:
+        caches = pd.DataFrame(caches)
+        caches['last_modified'] = pd.to_datetime(caches['last_modified'], unit='us').dt.floor('s')
+        caches = caches[['cache_key', 'namespace', 'backend', 'state', 'last_modified', 'user', 'commit_hash', 'path']]
+        caches = caches.to_string()
+
+
+    if pager:
+        click.echo_via_pager(caches)
+    else:
+        click.echo(caches)
+
+
+@cli.command('delete')
+@click.argument('cache_key')
+@click.option('--namespace', help='Namespace for the cache')
+@click.option('--backend', help='Backend to use')
+@click.option('--location', help='Location to search', default='root')
+@click.option('--force', '-f', is_flag=True, help='Force deletion without confirmation')
+@click.option('--metadata-only', '-m', is_flag=True, help='Only delete the metadata for the cache, not the underlying data.')
+def delete_cache(cache_key, namespace, backend, location, force, metadata_only):
+    """Clear cache entries."""
+    caches = list_helper(cache_key, namespace, backend, location)
+    config = get_config(location=location, requested_parameters=Cache.config_parameters)
+
+    click.confirm(f"Are you sure you want to delete {len(caches)} cache entries?", abort=True)
+
+    for cache in caches:
+        cache = Cache(config, cache.cache_key, cache.namespace, None, location, cache.backend, {})
+        click.echo(f"Deleting {cache.cache_key} from {cache.location} with backend {cache.backend_name}.")
+        if metadata_only:
+            cache._delete_metadata()
+        else:
+            cache.delete()
+
+
+@cli.command('print-config')
+@click.option('--location', help='Location to search', default='root')
+@click.option('--backend', help='Backend to use')
+@click.option('--show-secrets', '-s', is_flag=True, help='Only delete the metadata for the cache, not the underlying data.')
+def get_config_value(location, backend, show_secrets):
+    """Get configuration value for a specific key."""
+    if backend:
+        backend_classes = [get_backend_by_name(backend)]
+    else:
+        backend_classes = [Cache] + list(registered_backends.values())
+
+    for backend_class in backend_classes:
+        if show_secrets:
+            config = get_config(location=location, requested_parameters=backend_class.config_parameters,
+                                backend_name=backend_class.backend_name, mask_secrets=False)
+        else:
+            config = get_config(location=location, requested_parameters=backend_class.config_parameters,
+                                backend_name=backend_class.backend_name, mask_secrets=True)
+
+        click.echo(backend_class.backend_name.title())
+        for key, value in config.items():
+            click.echo(f"\t{key}: {value}")
+        click.echo()
+
+def main():
+    return cli()
+
+
+if __name__ == '__main__':
+    main()

nuthatch/config.py

@@ -0,0 +1,94 @@
+"""The config module is used to get the config for a given location and backend."""
+from pathlib import Path
+import os
+import tomllib
+
+dynamic_parameters = {}
+
+def config_parameter(parameter_name, location='root', backend=None, secret=False):
+    """A decorator to register a function as a dynamic parameter.
+
+    Args:
+        parameter_name (str): The name of the parameter.
+        location (str, optional): The location to register the parameter.
+        backend (str, optional): The backend to register the parameter.
+    """
+    def decorator(function):
+        if location not in dynamic_parameters:
+            dynamic_parameters[location] = {}
+        if backend and backend not in dynamic_parameters[location]:
+            dynamic_parameters[location][backend] = {}
+
+        if backend:
+            dynamic_parameters[location][backend][parameter_name] = (function, secret)
+        else:
+            dynamic_parameters[location][parameter_name] = (function, secret)
+    return decorator
+
+def _is_fs_root(p):
+    """Check if a path is the root of a filesystem."""
+    return os.path.splitdrive(str(p))[1] == os.sep
+
+def get_config(location='root', requested_parameters=[], backend_name=None, mask_secrets=False):
+    """Get the config for a given location and backend.
+
+    Args:
+        location (str, optional): The location to get the config for.
+        requested_parameters (list, optional): The parameters to get the config for.
+        backend_name (str, optional): The backend to get the config for.
+    """
+    #Find pyproject.toml or nuthatch.ini
+    current_directory = Path.cwd()
+
+    config_file = None
+    while not _is_fs_root(current_directory):
+        if current_directory.joinpath('pyproject.toml').exists():
+            config_file = current_directory.joinpath('pyproject.toml')
+
+        current_directory = current_directory.parent
+
+    #TODO: enable ini and environment variable configuration
+
+    with open(config_file, "rb") as f:
+        config = tomllib.load(f)
+
+    # If it's root allow the base parameters to be used and root to be set
+    location_params = {}
+    if location == 'root':
+        location_params = config['tool']['nuthatch']
+        if location in config['tool']['nuthatch']:
+            location_params.update(config['tool']['nuthatch'][location])
+    else:
+        if location in config['tool']['nuthatch']:
+            location_params = config['tool']['nuthatch'][location]
+
+    if backend_name and backend_name in location_params:
+        backend_specific_params = config['tool']['nuthatch'][location][backend_name]
+    else:
+        backend_specific_params = {}
+
+    # Merge the two together
+    merged_config = backend_specific_params | location_params
+
+    filtered_config = {k: merged_config[k] for k in merged_config if k in requested_parameters}
+
+    # Now call all the relevant config registrations and add them
+    for p in requested_parameters:
+        if location in dynamic_parameters:
+            if backend_name in dynamic_parameters[location] and p in dynamic_parameters[location][backend_name]:
+                secret = dynamic_parameters[location][backend_name][p][1]
+                param = dynamic_parameters[location][backend_name][p][0]()
+                if secret and mask_secrets:
+                    filtered_config[p] = '*'*len(param)
+                else:
+                    filtered_config[p] = param
+            elif p in dynamic_parameters[location]:
+                secret = dynamic_parameters[location][p][1]
+                param = dynamic_parameters[location][p][0]()
+                if secret and mask_secrets:
+                    filtered_config[p] = '*'*len(param)
+                else:
+                    filtered_config[p] = param
+
+
+    return filtered_config

nuthatch/memoizer.py

@@ -0,0 +1,67 @@
+"""The memoizer is a module that memoizes data in memory."""
+import sys
+import xarray as xr
+import dask.dataframe as dd
+from .config import get_config
+
+memoized_objects = {}
+cache_key_lru = []
+
+def save_to_memory(cache_key, data):
+    """Save data to memory.
+
+    Implements special handling for xarray and dask dataframes.
+    Evicts the least recently used object when the memory limit is reached.
+
+    Args:
+        cache_key (str): The key to save the data to.
+        data (any): The data to save to memory.
+    """
+    if isinstance(data, xr.Dataset):
+        data = data.persist()
+    elif isinstance(data, dd.DataFrame):
+        data = data.persist()
+    else:
+        pass
+
+    max_size = get_config(location='root', requested_parameters=['maximum_memory_usage'])
+    if 'maximum_memory_usage' in max_size:
+        max_size = max_size['maximum_memory_usage']
+    else:
+        max_size = 100*10^6
+
+    if(sys.getsizeof(data) > max_size):
+        print("WARNING: Data too large to memoize.")
+
+    while(sys.getsizeof(memoized_objects) + sys.getsizeof(data) > max_size):
+        del memoized_objects[cache_key_lru[0]]
+        del cache_key_lru[0]
+
+    memoized_objects[cache_key] = data
+
+    if cache_key in cache_key_lru:
+        cache_key_lru.remove(cache_key)
+
+    cache_key_lru.append(cache_key)
+
+
+def recall_from_memory(cache_key):
+    """Recall data from memory.
+
+    Refreshes the LRU cache when the data is recalled.
+
+    Args:
+        cache_key (str): The key to recall the data from.
+
+    Returns:
+        The memoized object.
+    """
+    if cache_key in memoized_objects:
+        # refresh the lru
+        cache_key_lru.remove(cache_key)
+        cache_key_lru.append(cache_key)
+
+        # return the object
+        return memoized_objects[cache_key]
+    else:
+        return None