Source code for sparc.download_data

"""Download the pseudopotential and other related files after sparc-x-api is installed

Run:

python -m sparc.download_data
"""

import hashlib
import shutil
import tempfile
import zipfile
from io import BytesIO
from pathlib import Path

# import urllib.request
from urllib.request import urlopen

from .common import psp_dir

sparc_tag = "b702c1061400a2d23c0e223e32182609d7958156"
sparc_source_url = "https://github.com/SPARC-X/SPARC/archive/{sparc_tag}.zip"
# This is a all_psp8_checksum
all_psp8_checksum = "5ef42c4a81733a90b0e080b771c5a73a"


[docs] def download_psp(sparc_tag=sparc_tag, psp_dir=psp_dir): """Download the external PSPs into the sparc/psp folder Arguments: sparc_tag (str): Commit hash or git tag for the psp files psp_dir (str or PosixPath): Directory to download the psp files """ if is_psp_download_complete(): print("PSPs have been successfully downloaded!") return download_url = sparc_source_url.format(sparc_tag=sparc_tag) print(f"Download link: {download_url}") with tempfile.TemporaryDirectory() as tmpdir: tmpdir = Path(tmpdir) with urlopen(download_url) as zipresp: with zipfile.ZipFile(BytesIO(zipresp.read())) as zfile: zfile.extractall(tmpdir) # print(list(os.walk(tmpdir))) source_dir = next(tmpdir.glob("SPARC-*/psps")) print(f"Found source_dir at {source_dir}") if not source_dir.is_dir(): raise FileNotFoundError("Error downloading or extracting zip") print(f"Moving psp files to {psp_dir}") for ext in ("*.psp8", "*.psp", "*.pot"): for pspf in source_dir.glob(ext): print(f"Found {pspf} --> {psp_dir}") shutil.copy(pspf, psp_dir) if not is_psp_download_complete(psp_dir): raise RuntimeError(f"Files downloaded to {psp_dir} have different checksums!") return
[docs] def checksum_all(psp_dir=psp_dir, extension="*.psp8"): """Checksum all the files under the psp_dir to make sure the psp8 files are the same as intended Arguments: psp_dir (str or PosixPath): Directory for the psp files extension (str): Search pattern for the psp files, either '.psp', '.psp8' or '.pot' Returns: str: Checksum for all the files concatenated """ checker = hashlib.md5() psp_dir = Path(psp_dir) # Use sorted to make sure file order is correct for filename in sorted(psp_dir.glob(extension)): # Open the file in binary mode and update the group checksum print(f"Checking {filename}") with open(filename, "r") as f: f_checker = hashlib.md5() content = f.read().encode("utf8") f_checker.update(content) checker.update(f_checker.hexdigest().encode("ascii")) final_checksum = checker.hexdigest() print(f"Final checksum is {final_checksum}") return final_checksum
[docs] def is_psp_download_complete(psp_dir=psp_dir): return checksum_all(psp_dir) == all_psp8_checksum
if __name__ == "__main__": print("Running command-line psp downloader") download_psp()