"""Utilities for downloading and unpacking test data packages
WARNING: This module must not depend on a working pynbody installation.
It is used during CI to download test data before pynbody is built/installed.
Only use standard library modules and osfclient (which is available in CI).
"""
import argparse
import os
import pathlib
import shutil
import ssl
import tarfile
import urllib.request
import osfclient
# Cache OSF file objects to avoid repeating http requests
_OSF_STORAGE_CACHE = {}
test_data_packages = {
'swift': {'verify_path': 'SWIFT',
'archive_name': 'swift.tar.gz'},
'swift_isolated': {'verify_path': 'SWIFT/isolated_0008.hdf5',
'extract_path': 'SWIFT',
'archive_name': 'swift_isolated.tar.gz'},
'swift_planetary': {'verify_path': 'SWIFT/planetary.hdf5',
'extract_path': 'SWIFT',
'archive_name': 'swift_planetary.tar.gz'},
'adaptahop_longint': {'verify_path': 'adaptahop_longint',
'archive_name': 'adaptahop_longint.tar.gz'},
'arepo': {'verify_path': 'arepo',
'archive_name': 'arepo.tar.gz'},
'gadget': {'verify_path': 'gadget2',
'archive_name': 'gadget.tar.gz'},
'hbt': {'verify_path': 'gadget4_subfind_HBT',
'archive_name': 'gadget4_subfind_HBT.tar.gz'},
'gasoline_ahf': {'verify_path': 'gasoline_ahf',
'archive_name': 'gasoline.tar.gz'},
'gizmo': {'verify_path': 'gizmo',
'archive_name': 'gizmo.tar.gz'},
'grafic': {'verify_path': 'grafic_test',
'archive_name': 'grafic.tar.gz'},
'lpicola': {'verify_path': 'lpicola',
'archive_name': 'lpicola.tar.gz'},
'nchilada': {'verify_path': 'nchilada_test',
'archive_name': 'nchilada.tar.gz'},
'ramses': {'verify_path': 'ramses',
'archive_name': 'ramses.tar.gz'},
'rockstar': {'verify_path': 'rockstar',
'archive_name': 'rockstar.tar.gz'},
'subfind': {'verify_path': 'subfind',
'archive_name': 'subfind.tar.gz'},
'tng_subfind': {'verify_path': 'arepo/tng',
'archive_name': 'tng_subfind.tar.gz'},
'pkdgrav3': {'verify_path': 'pkdgrav3',
'archive_name': 'pkdgrav3.tar.gz'},
'tutorial_gadget4': {'verify_path': 'tutorial_gadget4/subhalo_desc_035.hdf5',
'archive_name': 'tutorial_gadget4.tar.gz',
'extended': True},
'tutorial_gadget': {'verify_path': 'tutorial_gadget/snapshot_020',
'archive_name': 'tutorial_gadget.tar.gz',
'extended': True},
'tiny_FIRE' : {'verify_path': 'tiny_FIRE',
'archive_name': 'tiny_FIRE.tar.gz'}
}
osf_project_id = '5m6zs'
[docs]
def precache_test_data(verbose=False, extended=False):
"""Download and unpack all test data packages.
By default, only the packages required for resting are downloaded. To include extended packages for documentation/
tutorials, set extended=True.
"""
for package_name, package in test_data_packages.items():
# only download packages that are not marked as 'extended' packages - unless extended=True
if extended or not package.get('extended', False):
_download_and_unpack_test_data_if_not_present(package, package_name, verbose)
[docs]
def test_data_hash():
"""Return a hash representing the data packages to be downloaded"""
# print a hex digest of the hash of the test data package urls
import hashlib
m = hashlib.sha256()
for package_name in test_data_packages:
m.update(test_data_packages[package_name]['archive_name'].encode())
m.update(osf_project_id.encode())
return m.hexdigest()
[docs]
def get_osf_file_object(osf_project_id, archive_name):
"""Retrieve the OSF file object for the given archive name in the specified OSF project."""
# Open the OSF project storage, if we didn't already
if osf_project_id not in _OSF_STORAGE_CACHE:
osf = osfclient.OSF()
osf_project = osf.project(osf_project_id)
osf_storage = osf_project.storage('osfstorage')
_OSF_STORAGE_CACHE[osf_project_id] = {}
for file in osf_storage.files:
_OSF_STORAGE_CACHE[osf_project_id][file.name] = file
# Return the cached file
return _OSF_STORAGE_CACHE[osf_project_id][archive_name]
[docs]
def download_and_unpack_test_data(archive_name, unpack_path="", verbose=False):
"""Download and unpack test data with the given archive name and unpack path.
Equivalent to running:
wget https://zenodo.org/record/.../files/{archive_name}?download=1
tar -xzf {archive_name}
"""
osf_file = get_osf_file_object(osf_project_id, archive_name)
unpack_path = f"testdata/{unpack_path}"
if not os.path.exists(unpack_path):
os.makedirs(unpack_path, exist_ok=True)
# Download to a temporary file first
temp_file = f"{archive_name}.tmp"
try:
if verbose:
print(f"Downloading {archive_name}")
with open(temp_file, 'wb') as f:
osf_file.write_to(f)
if verbose:
print(f"Extracting {archive_name} to {unpack_path}")
# Extract from the downloaded file
with tarfile.open(temp_file) as tar:
tar.extractall(unpack_path, filter='data')
if verbose:
print(f"Successfully unpacked {archive_name}")
finally:
# Clean up temporary file
if os.path.exists(temp_file):
os.remove(temp_file)
[docs]
def ensure_test_data_available(*package_names):
"""Ensure that the specified test data packages are available in the testdata directory."""
for package_name in package_names:
if package_name not in test_data_packages:
raise ValueError(f"Test data package {package_name} not found in test_data_packages")
package = test_data_packages[package_name]
_download_and_unpack_test_data_if_not_present(package, package_name, False)
def _download_and_unpack_test_data_if_not_present(package, package_name, verbose=False):
if not pathlib.Path(f"testdata/{package['verify_path']}").exists():
if verbose:
print(f"Test data package '{package_name}' not found, downloading...")
download_and_unpack_test_data(package['archive_name'], package.get('extract_path', ''), verbose)
elif verbose:
print(f"Test data package '{package_name}' already exists, skipping")
_ionfrac_tables_osf_project_id = "z46rq"
_required_ionfrac_tables = ("hm12", "fg20") # Tables needed by the tests
[docs]
def download_ionfrac_table(name, destination):
"""Download an ion fraction table from the pynbody data repository
Note that this is also called from pynbody.analysis.ionfrac to fetch
the tables when they are needed and not already present.
"""
osf_file = get_osf_file_object(_ionfrac_tables_osf_project_id, f'{name}.npz')
with open(destination, 'wb') as f:
osf_file.write_to(f)
[docs]
def precache_ionfrac_tables(verbose=False):
"""Download all ionfrac tables needed for tests"""
os.makedirs("ionfrac_tables", exist_ok=True)
for name in _required_ionfrac_tables:
if verbose:
print(f"Downloading table: {name}")
download_ionfrac_table(name, pathlib.Path("ionfrac_tables") / pathlib.Path(name+".npz"))
[docs]
def ionfrac_tables_hash():
"""Return a hash of the ionfrac table names"""
import hashlib
m = hashlib.sha256()
for name in _required_ionfrac_tables:
m.update(name.encode())
return m.hexdigest()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Download data files needed for tests')
group = parser.add_mutually_exclusive_group()
group.add_argument("--fetch-testdata", action="store_true", help="Download pynbody test data files")
group.add_argument("--hash-testdata", action="store_true", help="Output the hash of the test data files")
group.add_argument("--fetch-ionfrac", action="store_true", help="Download ionfrac table files needed for tests")
group.add_argument("--hash-ionfrac", action="store_true", help="Output the hash of the ionfrac table names")
args = parser.parse_args()
if args.fetch_testdata:
precache_test_data(verbose=True)
elif args.hash_testdata:
print(test_data_hash())
elif args.fetch_ionfrac:
precache_ionfrac_tables(verbose=True)
elif args.hash_ionfrac:
print(ionfrac_tables_hash())
else:
raise RuntimeError("No operation specified!")
