_utils.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. import os
  2. import shutil
  3. from ._registry import method_files_map
  4. from scipy._lib._array_api import xp_capabilities
  5. try:
  6. import platformdirs
  7. except ImportError:
  8. platformdirs = None # type: ignore[assignment]
  9. def _clear_cache(datasets, cache_dir=None, method_map=None):
  10. if method_map is None:
  11. # Use SciPy Datasets method map
  12. method_map = method_files_map
  13. if cache_dir is None:
  14. # Use default cache_dir path
  15. if platformdirs is None:
  16. # platformdirs is pooch dependency
  17. raise ImportError("Missing optional dependency 'pooch' required "
  18. "for scipy.datasets module. Please use pip or "
  19. "conda to install 'pooch'.")
  20. cache_dir = platformdirs.user_cache_dir("scipy-data")
  21. if not os.path.exists(cache_dir):
  22. print(f"Cache Directory {cache_dir} doesn't exist. Nothing to clear.")
  23. return
  24. if datasets is None:
  25. print(f"Cleaning the cache directory {cache_dir}!")
  26. shutil.rmtree(cache_dir)
  27. else:
  28. if not isinstance(datasets, list | tuple):
  29. # single dataset method passed should be converted to list
  30. datasets = [datasets, ]
  31. for dataset in datasets:
  32. assert callable(dataset)
  33. dataset_name = dataset.__name__ # Name of the dataset method
  34. if dataset_name not in method_map:
  35. raise ValueError(f"Dataset method {dataset_name} doesn't "
  36. "exist. Please check if the passed dataset "
  37. "is a subset of the following dataset "
  38. f"methods: {list(method_map.keys())}")
  39. data_files = method_map[dataset_name]
  40. data_filepaths = [os.path.join(cache_dir, file)
  41. for file in data_files]
  42. for data_filepath in data_filepaths:
  43. if os.path.exists(data_filepath):
  44. print("Cleaning the file "
  45. f"{os.path.split(data_filepath)[1]} "
  46. f"for dataset {dataset_name}")
  47. os.remove(data_filepath)
  48. else:
  49. print(f"Path {data_filepath} doesn't exist. "
  50. "Nothing to clear.")
  51. @xp_capabilities(out_of_scope=True)
  52. def clear_cache(datasets=None):
  53. """
  54. Cleans the scipy datasets cache directory.
  55. If a scipy.datasets method or a list/tuple of the same is
  56. provided, then clear_cache removes all the data files
  57. associated to the passed dataset method callable(s).
  58. By default, it removes all the cached data files.
  59. Parameters
  60. ----------
  61. datasets : callable or list/tuple of callable or None
  62. Examples
  63. --------
  64. >>> from scipy import datasets
  65. >>> ascent_array = datasets.ascent()
  66. >>> ascent_array.shape
  67. (512, 512)
  68. >>> datasets.clear_cache([datasets.ascent])
  69. Cleaning the file ascent.dat for dataset ascent
  70. """
  71. _clear_cache(datasets)