airoh.acquisition

  1# src/airoh/acquisition.py
  2"""Tasks for acquiring data assets: download, symlink/copy, or git submodule.
  3
  4These are the dependency-light acquisition tasks. Datalad-backed acquisition
  5lives in `airoh.datalad`, which is gated behind the optional `datalad` extra.
  6"""
  7import os
  8import shutil
  9from pathlib import Path
 10from urllib.request import Request, urlopen
 11
 12from invoke import task
 13
 14
 15@task(
 16    help={
 17        "name": "Logical name of the file, as defined in the 'files' section of invoke.yaml."
 18    }
 19)
 20def download_data(c, name):
 21    """🌐 Download a single file from a URL using urllib.
 22
 23    Looks up the file entry by logical name in invoke.yaml under the `files`
 24    key, then downloads it to the configured output path. Skips the download
 25    if the output file already exists and is non-empty. Uses a `.part` temp
 26    file during transfer and atomically replaces the target on success.
 27
 28    Parameters
 29    ----------
 30    c : invoke.Context
 31        The Invoke context.
 32    name : str
 33        Logical name of the file to download, matching a key under `files`
 34        in invoke.yaml. Each entry must define `url` and `output_file`.
 35
 36    Raises
 37    ------
 38    ValueError
 39        If `name` is not found under `files` in invoke.yaml, or if the
 40        matched entry is missing `url` or `output_file`.
 41    RuntimeError
 42        If the download completes with 0 bytes, or if any network error occurs.
 43
 44    Examples
 45    --------
 46    ```bash
 47    inv acquisition.download-data --name my_dataset
 48    ```
 49    """
 50    files = c.config.get("files", {})
 51    if name not in files:
 52        raise ValueError(f"❌ No file config found for '{name}' in invoke.yaml.")
 53
 54    entry = files[name]
 55    url = entry.get("url")
 56    output_file = entry.get("output_file")
 57
 58    if not url or not output_file:
 59        raise ValueError(
 60            f"❌ Entry for '{name}' must define both 'url' and 'output_file'."
 61        )
 62
 63    output_path = Path(output_file)
 64    tmp_path = output_path.with_suffix(output_path.suffix + ".part")
 65
 66    if output_path.exists() and output_path.stat().st_size > 0:
 67        print(f"🫧 Skipping {name}: {output_file} already exists.")
 68        return
 69
 70    output_path.parent.mkdir(parents=True, exist_ok=True)
 71    tmp_path.unlink(missing_ok=True)
 72
 73    print(f"📥 Downloading '{name}' from {url}")
 74    print(f"📁 Target: {output_file}")
 75
 76    req = Request(
 77        url,
 78        headers={
 79            "User-Agent": "Mozilla/5.0",
 80            "Accept": "*/*",
 81        },
 82    )
 83
 84    try:
 85        with urlopen(req, timeout=60) as response, tmp_path.open("wb") as f:
 86            total = 0
 87            while True:
 88                chunk = response.read(8192)
 89                if not chunk:
 90                    break
 91                f.write(chunk)
 92                total += len(chunk)
 93
 94        if total == 0:
 95            tmp_path.unlink(missing_ok=True)
 96            raise RuntimeError(f"❌ Downloaded 0 bytes for '{name}'.")
 97
 98        tmp_path.replace(output_path)
 99
100    except Exception as e:
101        tmp_path.unlink(missing_ok=True)
102        raise RuntimeError(f"❌ Failed to download '{name}' from {url}: {e}") from e
103
104    print(f"✅ Downloaded {name} to {output_file} ({output_path.stat().st_size} bytes)")
105
106
107@task(
108    help={
109        "name": "Logical name of the asset, as defined in the 'files' section of invoke.yaml.",
110        "source": (
111            "Path to already-present data to link instead of downloading. "
112            "Overrides the entry's optional 'source' key."
113        ),
114        "copy": "Copy the source data instead of symlinking it.",
115    }
116)
117def fetch_data(c, name, source=None, copy=False):
118    """📦 Make a data asset available: symlink existing data, or download it.
119
120    Looks up the asset entry by logical name in invoke.yaml under the `files`
121    key. Decides where the data comes from in this order:
122
123    1. the `source` argument (e.g. ``invoke ... --source /path``),
124    2. the entry's optional `source` key in invoke.yaml,
125    3. otherwise, the entry's `url`, downloaded via `download_data`.
126
127    When a source path is found, it is symlinked (or copied, with ``copy=True``)
128    to the entry's `output_file`. Both single files and whole directories are
129    supported. The operation is idempotent: a link that already points at the
130    source is left untouched.
131
132    Parameters
133    ----------
134    c : invoke.Context
135        The Invoke context.
136    name : str
137        Logical name of the asset, matching a key under `files` in invoke.yaml.
138        Each entry must define `output_file`, plus `url` and/or `source`.
139    source : str, optional
140        Path to existing data to link/copy. Takes precedence over the entry's
141        `source` key. If neither is set, the asset is downloaded from `url`.
142    copy : bool, optional
143        Copy the source data instead of symlinking it (default: False).
144
145    Raises
146    ------
147    ValueError
148        If `name` is not found under `files`, if the entry has no `output_file`,
149        if the resolved source path does not exist, or if `output_file` already
150        exists as something other than a symlink.
151
152    Examples
153    --------
154    ```bash
155    inv acquisition.fetch-data --name my_dataset                       # download
156    inv acquisition.fetch-data --name my_dataset --source /data/mine   # symlink
157    inv acquisition.fetch-data --name my_dataset --source /data/mine --copy
158    ```
159    """
160    files = c.config.get("files", {})
161    if name not in files:
162        raise ValueError(f"❌ No file config found for '{name}' in invoke.yaml.")
163
164    entry = files[name]
165    output_file = entry.get("output_file")
166    if not output_file:
167        raise ValueError(f"❌ Entry for '{name}' must define 'output_file'.")
168
169    resolved_source = source or entry.get("source")
170    if not resolved_source:
171        download_data(c, name)
172        return
173
174    _link_data(name, resolved_source, output_file, copy=copy)
175
176
177def _link_data(name, source, output_file, copy=False):
178    """Symlink (or copy) existing data at `source` to `output_file`.
179
180    Idempotent: if `output_file` is already a symlink pointing at `source`, it is
181    left untouched; a symlink pointing elsewhere is repointed. A pre-existing
182    non-symlink at `output_file` raises, to avoid clobbering real data.
183    """
184    source_path = Path(source).expanduser().resolve()
185    output_path = Path(output_file)
186
187    if not source_path.exists():
188        raise ValueError(f"❌ Source for '{name}' does not exist: {source_path}")
189
190    if output_path.is_symlink():
191        if output_path.resolve() == source_path:
192            print(f"🫧 Skipping {name}: {output_file} already links to {source_path}.")
193            return
194        output_path.unlink()
195    elif output_path.exists():
196        raise ValueError(
197            f"❌ {output_file} already exists and is not a symlink; remove it or "
198            f"run `invoke clean` before linking '{name}'."
199        )
200
201    output_path.parent.mkdir(parents=True, exist_ok=True)
202
203    if copy:
204        copier = shutil.copytree if source_path.is_dir() else shutil.copy2
205        copier(source_path, output_path)
206        print(f"✅ Copied {name} from {source_path} to {output_file}.")
207    else:
208        output_path.symlink_to(source_path)
209        print(f"✅ Linked {name}: {output_file} -> {source_path}.")
210
211
212@task
213def ensure_submodule(c, path, recursive=True):
214    """🔄 Ensure a git submodule is initialized and up to date.
215
216    Initializes or updates a specified git submodule recursively.
217
218    Parameters
219    ----------
220    c : invoke.Context
221        The Invoke context.
222    path : str
223        Path to the submodule directory.
224    recursive : boolean, default True
225        Update all submodules recursively
226
227    Examples
228    --------
229    ```bash
230    inv acquisition.ensure-submodule --path src/external-lib
231    ```
232    """
233    if recursive:
234        flag_recursive = "--recursive"
235    else:
236        flag_recursive = ""
237
238    if not os.path.exists(path) or not os.path.exists(os.path.join(path, ".git")):
239        print(f"📦 Initializing submodule at {path}...")
240        c.run(f"git submodule update --init {flag_recursive} {path}")
241    else:
242        print(f"🔄 Updating submodule at {path}...")
243        c.run(f"git submodule update --remote {path}")