diff --git a/src/hr_selection_function/data.py b/src/hr_selection_function/data.py index d7e0575..85d487e 100644 --- a/src/hr_selection_function/data.py +++ b/src/hr_selection_function/data.py @@ -3,6 +3,7 @@ from hr_selection_function.config import _CONFIG from tqdm import tqdm from pathlib import Path from functools import wraps +import shutil def set_data_directory(directory: Path | str): @@ -27,6 +28,10 @@ def download_data(redownload: bool = False): return _download_file(_CONFIG["data_url"], unzip=True, write_location=_CONFIG["data_dir"]) + if not _check_data_downloaded(): + raise RuntimeError( + "Data download failed. Data is not as expected, or where it should be." + ) _CONFIG["data_already_downloaded"] = True @@ -34,10 +39,12 @@ def requires_data(func): """Wrapper for some function func that ensures that the data directory has been created before proceeding. """ + @wraps(func) def inner(*args, **kwargs): download_data(redownload=False) func(*args, **kwargs) + return inner @@ -51,35 +58,39 @@ def _download_file( write_location = _CONFIG["data_dir"] write_location = Path(write_location) - print(f"Downloading data for hr_selection_function to {write_location}") - if unzip: write_location = write_location / "_temp.zip" + print(f"Downloading data for hr_selection_function to {write_location}") + # Fetch initial dataset response = requests.get(data_link, stream=True) - with tqdm.wrapattr( - open(write_location, "wb"), - "write", - unit="GB", - unit_scale=True, - unit_divisor=1024**3, - miniters=1, - desc="Data", - total=int(response.headers.get("content-length", 0)), - ) as file: - for chunk in response.iter_content(chunk_size=4096): - file.write(chunk) + with open(write_location, "wb") as handle: + with tqdm.wrapattr( + handle, + "write", + unit="GB", + unit_scale=True, + unit_divisor=1024**3, + miniters=1, + desc="Data", + total=int(response.headers.get("content-length", 0)), + ) as file: + for chunk in response.iter_content(chunk_size=4096): + file.write(chunk) if not unzip: return # Unzip data - print("Unzipping data...") - # Todo + print(f"Unzipping data to {write_location}") + shutil.unpack_archive(write_location, write_location.parent) # Delete raw .zip file - print("Removing .zip file...") + print(f"Removing .zip file at {write_location}") + write_location.unlink() + + print("Data downloaded!") def _check_data_directory(directory: Path | str) -> tuple[Path, bool]: @@ -107,12 +118,14 @@ def _check_data_directory(directory: Path | str) -> tuple[Path, bool]: f"{directory}..." ) directory.mkdir(parents=True) - first_time = False + first_time = True return directory, first_time def _check_data_downloaded(directory: Path | None = None) -> bool: + # Todo this function should really check file size too + if directory is None: directory = _CONFIG["data_dir"] directory = Path(directory) diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..4411299 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,139 @@ +"""Test data handling and download in the library.""" + +from hr_selection_function.data import ( + _download_file, + _check_data_directory, + _check_data_downloaded, + set_data_directory, +) +from hr_selection_function.config import _CONFIG, _DEFAULT_DIRECTORY +from pathlib import Path +import shutil +import warnings +import pytest + + +location = Path.cwd() / "hr_selection_function_test_data_578322146558" + + +def test__download_file(): + file = "https://sample-files.com/downloads/compressed/zip/basic-text.zip" # Todo find something better + location.mkdir(exist_ok=True, parents=True) + + try: + _download_file(file, unzip=True, write_location=location) + + # Check all is as expected + assert location.exists() + assert (location / "file1.txt").exists() + assert (location / "file2.txt").exists() + assert (location / "file3.txt").exists() + assert (location / "file4.txt").exists() + assert (location / "file5.txt").exists() + assert not (location / "_temp.zip").exists() + + finally: + shutil.rmtree(location) + + +def test_file_download_no_zip(): + file = "https://sample-files.com/downloads/compressed/zip/basic-text.zip" # Todo find something better + location.mkdir(exist_ok=True, parents=True) + + write = location / "a_zip_file.zip" + + try: + _download_file(file, unzip=False, write_location=write) + + # Check all is as expected + assert location.exists() + assert write.exists() + + finally: + shutil.rmtree(location) + + +def test__check_data_directory(): + # Assert arg type + with pytest.raises( + ValueError, match="Data directory path must be a pathlib.Path or string." + ): + _check_data_directory(42) + + # Check directory is made in usual circumstance + try: + location_made, first_time = _check_data_directory(location) + assert first_time + assert location.exists() + assert location == location_made + finally: + location.rmdir() + + # Check again, if it's a string is made in usual circumstance + try: + location_made, first_time = _check_data_directory(str(location)) + assert first_time + assert location.exists() + assert location == location_made + finally: + location.rmdir() + + # Check first time arg + # Check directory is made in usual circumstance + location.mkdir(parents=True) + try: + location_made, first_time = _check_data_directory(location) + assert first_time is False + assert location.exists() + finally: + location.rmdir() + + +def test_download_data(): + warnings.warn("Todo once data is online somewhere") + + +def test__check_data_downloaded(): + # Initial setup + (location / "nstars_models").mkdir(parents=True) + + files = [ + "density_hp7.parquet", + "mcmc_samples.parquet", + "subsample_cuts_hp7.parquet", + ] + files = files + [f"nstars_models/{i}.ubj" for i in range(250)] + files = [location / file for file in files] + for file in files: + file.touch() + + try: + assert _check_data_downloaded(location) + + files[0].unlink() + assert _check_data_downloaded(location) is False + + files[0].touch() + assert _check_data_downloaded(location) + + files[-100].unlink() # Checks sensitivity to n_stars models + assert _check_data_downloaded(location) is False + + finally: + shutil.rmtree(location) + + +def test_set_data_directory(): + # Tests won't work if this isn't right + assert _CONFIG["data_dir"] == _DEFAULT_DIRECTORY + assert not location.exists() + + try: + set_data_directory(location) + assert location.exists() + assert _CONFIG["data_dir"] == location + assert _CONFIG["first_run"] + assert _CONFIG["data_already_downloaded"] is False + + finally: + shutil.rmtree(location) diff --git a/uv.lock b/uv.lock index 32680b4..95388e2 100644 --- a/uv.lock +++ b/uv.lock @@ -985,7 +985,7 @@ requires-dist = [ { name = "pyarrow" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "ruff", marker = "extra == 'dev'" }, - { name = "tqdm", specifier = ">=4.67.1" }, + { name = "tqdm", specifier = ">=3.0.0" }, { name = "xgboost", specifier = ">=3.0.0" }, ] provides-extras = ["dev", "notebooks"]