Justin Fisher
144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Download gridMET daily NetCDF files into year-based folders.
|
|
|
|
Default layout:
|
|
data/gridmet/1991/sph.nc
|
|
data/gridmet/1991/rmax.nc
|
|
data/gridmet/1991/rmin.nc
|
|
|
|
gridMET source:
|
|
https://www.northwestknowledge.net/metdata/data/{variable}_{year}.nc
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
import tempfile
|
|
import urllib.error
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
|
|
DEFAULT_BASE_URL = "https://www.northwestknowledge.net/metdata/data"
|
|
DEFAULT_VARIABLES = ("sph", "rmax", "rmin")
|
|
|
|
|
|
def _parse_years(text: str) -> list[int]:
|
|
years: set[int] = set()
|
|
for part in text.split(","):
|
|
token = part.strip()
|
|
if not token:
|
|
continue
|
|
if "-" in token:
|
|
start_text, end_text = token.split("-", 1)
|
|
start = int(start_text)
|
|
end = int(end_text)
|
|
if start > end:
|
|
raise argparse.ArgumentTypeError(f"Invalid descending year range: {token}")
|
|
years.update(range(start, end + 1))
|
|
else:
|
|
years.add(int(token))
|
|
if not years:
|
|
raise argparse.ArgumentTypeError("At least one year is required.")
|
|
return sorted(years)
|
|
|
|
|
|
def _parse_variables(text: str) -> list[str]:
|
|
variables = [part.strip() for part in text.split(",") if part.strip()]
|
|
if not variables:
|
|
raise argparse.ArgumentTypeError("At least one variable is required.")
|
|
return variables
|
|
|
|
|
|
def _download_file(url: str, destination: Path, overwrite: bool) -> bool:
|
|
if destination.exists() and not overwrite:
|
|
print(f"skip existing {destination}")
|
|
return False
|
|
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.NamedTemporaryFile(
|
|
prefix=destination.stem + ".",
|
|
suffix=".download",
|
|
dir=destination.parent,
|
|
delete=False,
|
|
) as handle:
|
|
temp_path = Path(handle.name)
|
|
|
|
try:
|
|
print(f"download {url}")
|
|
with urllib.request.urlopen(url) as response, temp_path.open("wb") as output:
|
|
total_text = response.headers.get("Content-Length")
|
|
total = int(total_text) if total_text and total_text.isdigit() else None
|
|
copied = 0
|
|
while True:
|
|
chunk = response.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
output.write(chunk)
|
|
copied += len(chunk)
|
|
if total:
|
|
percent = copied / total * 100
|
|
print(f"\r {copied / 1024 / 1024:,.1f} MiB / {total / 1024 / 1024:,.1f} MiB ({percent:5.1f}%)", end="")
|
|
if total:
|
|
print()
|
|
|
|
temp_path.replace(destination)
|
|
print(f"saved {destination}")
|
|
return True
|
|
except (urllib.error.URLError, urllib.error.HTTPError, OSError) as exc:
|
|
temp_path.unlink(missing_ok=True)
|
|
print(f"failed {url}: {exc}", file=sys.stderr)
|
|
raise
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Download gridMET NetCDF files into data/gridmet/<year>/<variable>.nc."
|
|
)
|
|
parser.add_argument(
|
|
"--years",
|
|
type=_parse_years,
|
|
default=_parse_years("1991-2020"),
|
|
help="Comma-separated years/ranges. Default: 1991-2020.",
|
|
)
|
|
parser.add_argument(
|
|
"--variables",
|
|
type=_parse_variables,
|
|
default=list(DEFAULT_VARIABLES),
|
|
help="Comma-separated gridMET variables. Default: sph,rmax,rmin.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
type=Path,
|
|
default=Path("data/gridmet"),
|
|
help="Root folder for downloaded files. Default: data/gridmet.",
|
|
)
|
|
parser.add_argument(
|
|
"--base-url",
|
|
default=DEFAULT_BASE_URL,
|
|
help=f"Base gridMET download URL. Default: {DEFAULT_BASE_URL}.",
|
|
)
|
|
parser.add_argument(
|
|
"--overwrite",
|
|
action="store_true",
|
|
help="Redownload files that already exist.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
downloaded = 0
|
|
for year in args.years:
|
|
for variable in args.variables:
|
|
url = f"{args.base_url.rstrip('/')}/{variable}_{year}.nc"
|
|
destination = args.output_dir / str(year) / f"{variable}.nc"
|
|
if _download_file(url, destination, args.overwrite):
|
|
downloaded += 1
|
|
|
|
print(f"done; downloaded {downloaded} file(s)")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|