Parallel downloads into the same pooch cache

src/mccode_antlr/cli/cache.py

@@ -226,6 +226,7 @@ def warm_registries_via_pooch(registries) -> tuple[int, int]:
     -------
     (total_fetched, error_count)
     """
+    from mccode_antlr.reader.registry import pooch_fetch
     total = errors = 0
     seen: set = set()
 
@@ -243,7 +244,7 @@ def warm_registries_via_pooch(registries) -> tuple[int, int]:
         print(f"  [{reg.name}] downloading {len(files)} files …", flush=True)
         for fname in files:
             try:
-                p.fetch(fname)
+                pooch_fetch(p, fname)
                 total += 1
             except Exception as exc:
                 print(f"    WARNING: could not fetch {fname}: {exc}", flush=True)

src/mccode_antlr/reader/registry.py

@@ -61,6 +61,15 @@ def _dedupe_identical_paths(paths: list[Path]) -> Path | None:
     return sorted(paths, key=str)[0]
 
 
+def pooch_fetch(pooch_instance, filename: str) -> Path:
+    """pooch_instance.fetch(filename), safe for parallel downloads into one cache
+    (workaround until https://github.com/fatiando/pooch/pull/554 is released)."""
+    directory = getattr(pooch_instance, 'abspath', None)
+    if directory is not None:
+        Path(directory, filename).parent.mkdir(parents=True, exist_ok=True)
+    return Path(pooch_instance.fetch(filename))
+
+
 def _dedupe_identical_registry_entries(pooch_instance, names: list[str]) -> str | None:
     """Like _dedupe_identical_paths, but for a pooch-backed remote registry: pooch's
     registry index already stores a content hash per file (`pooch_instance.registry`
@@ -342,7 +351,7 @@ def is_available(self, name: str, ext: str = None, exact: bool = True):
         return self.pooch.registry_files(self.fullname(name, ext, exact))
 
     def path(self, name: str, ext: str = None, exact: bool = True) -> Path:
-        return Path(self.pooch.fetch(self.fullname(name, ext, exact)))
+        return pooch_fetch(self.pooch, self.fullname(name, ext, exact))
 
     def filenames(self) -> list[str]:
         return [x for x in self.pooch.registry_files]

tests/test_pooch_fetch_parallel.py

@@ -0,0 +1,50 @@
+"""Parallel downloads into one pooch cache (https://github.com/fatiando/pooch/issues/555)."""
+import functools
+import http.server
+import multiprocessing
+import threading
+
+import pooch
+
+from mccode_antlr.reader.registry import pooch_fetch
+
+NAMES = [f'd{i}/e{i}/f.txt' for i in range(100)]
+
+
+def _fetch_all(args):
+    port, cache = args
+    p = pooch.create(path=cache, base_url=f'http://127.0.0.1:{port}/',
+                     registry={name: None for name in NAMES})
+    failed = []
+    for name in NAMES:
+        try:
+            pooch_fetch(p, name)
+        except Exception as error:
+            failed.append(f'{name}: {error!r}')
+    return failed
+
+
+class _QuietHandler(http.server.SimpleHTTPRequestHandler):
+    def log_message(self, *args):
+        pass
+
+
+def test_parallel_fetch_into_new_directories(tmp_path):
+    served = tmp_path / 'served'
+    for name in NAMES:
+        (served / name).parent.mkdir(parents=True)
+        (served / name).write_text(name)
+    server = http.server.ThreadingHTTPServer(
+        ('127.0.0.1', 0), functools.partial(_QuietHandler, directory=str(served)))
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+    level = pooch.get_logger().level
+    pooch.get_logger().setLevel('WARNING')
+    try:
+        with multiprocessing.Pool(8) as pool:
+            results = pool.map(_fetch_all, [(server.server_address[1], str(tmp_path / 'cache'))] * 8)
+    finally:
+        pooch.get_logger().setLevel(level)
+        server.shutdown()
+    assert [f for failed in results for f in failed] == []
+    for name in NAMES:
+        assert (tmp_path / 'cache' / name).read_text() == name