datashare-python 0.11.2__py3-none-any.whl → 0.11.4__py3-none-any.whl

datashare_python/utils.py

@@ -10,6 +10,7 @@ import sys
 import threading
 import time
 from collections.abc import (
+    AsyncGenerator,
     AsyncIterable,
     Awaitable,
     Callable,
@@ -522,7 +523,7 @@ def _write_artifact_bytes(path: Path, artifact: bytes | BytesIO | Path) -> None:
 def _remove_previous_artifacts(path: Path) -> None:
     if isinstance(path, Path) and path.exists():
         with contextlib.suppress(FileNotFoundError):
-            path.rmdir()
+            shutil.rmtree(path)
     else:
         path.unlink(missing_ok=True)
 
@@ -766,59 +767,108 @@ class SharedResources:
         eviction_callback: Callable[[Any, Any], None] | None = None,
     ) -> None:
         self._eviction_callback = eviction_callback
-        self._cache = LRU(cache_size, self._eviction_callback)
+        self._cache = LRU(cache_size, self._on_evict)
         self._sentinel = object()
+        # Leased resources are tracked by id, evicting a leased resource defers the
+        # eviction callback until its last lease is released
+        self._lock = threading.Lock()
+        self._n_leases: dict[int, int] = dict()
+        self._pending_evictions: dict[int, tuple[Any, Any]] = dict()
 
     def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:  # noqa: ANN001
-        if self._eviction_callback is not None:
-            for k, v in self._cache.items():
-                self._eviction_callback(k, v)
+        for k, v in self._cache.items():
+            self._on_evict(k, v)
 
+    @contextlib.contextmanager
     def get_or_cache_resource(
         self, key: str, default_factory: Callable[[], Any]
-    ) -> Any:
-        value = self._cache.get(key, self._sentinel)
-        if value is not self._sentinel:
-            return value
-        # Get the value first to be sure to return the right one in case of concurrent
-        # access.
-        #
-        # Additionally, the factory can be long to complete (that's one of the reason
-        # we want to cache its results). Because we don't lock the factory call,
-        # together with the cache update, in case of concurrent access,
-        # quicker factories will complete first and will stay longer in the cache.
-        # This is fine, the alternative is to lock factory + cache update to avoid
-        # concurrent access but this will mean waiting for the first factory call to
-        # complete, this can potentially imply longer waits than no getting the value
-        # from the cache
-        value = default_factory()
-        self._cache[key] = value
-        return value
+    ) -> Generator[Any, None, None]:
+        """Get or cache a resource, leasing it for the duration of the context.
+
+        A resource evicted while leased stays usable, its eviction callback only runs
+        once all its leases have been released.
+
+        :param key: resource cache key
+        :param default_factory: resource factory, called on cache miss
+        :return: the cached resource
+        """
+        value = self._get_and_acquire(key)
+        if value is self._sentinel:
+            value = default_factory()
+            self._cache_and_acquire(key, value)
+        try:
+            yield value
+        finally:
+            self._release(value)
 
+    @contextlib.asynccontextmanager
     async def async_get_or_cache_resource(
         self, key: str, default_factory: Callable[[], Awaitable[Any]]
-    ) -> Any:
-        value = self._cache.get(key, self._sentinel)
-        if value is not self._sentinel:
-            return value
-        # Get the value first to be sure to return the right one in case of concurrent
-        # access.
-        #
-        # Additionally, the factory can be long to complete (that's one of the reason
-        # we want to cache its results). Because we don't lock the factory call,
-        # together with the cache update, in case of concurrent access,
-        # quicker factories will complete first and will stay longer in the cache.
-        # This is fine, the alternative is to lock factory + cache update to avoid
-        # concurrent access but this will mean waiting for the first factory call to
-        # complete, this can potentially imply longer waits than no getting the value
-        # from the cache
-        value = await default_factory()
-        self._cache[key] = value
+    ) -> AsyncGenerator[Any, None]:
+        """Async version of get_or_cache_resource.
+
+        :param key: resource cache key
+        :param default_factory: async resource factory, awaited on cache miss
+        :return: the cached resource
+        """
+        value = self._get_and_acquire(key)
+        if value is self._sentinel:
+            value = await default_factory()
+            self._cache_and_acquire(key, value)
+        try:
+            yield value
+        finally:
+            self._release(value)
+
+    def _get_and_acquire(self, key: str) -> Any:
+        with self._lock:
+            value = self._cache.get(key, self._sentinel)
+            if value is not self._sentinel:
+                self._acquire(value)
         return value
 
+    def _cache_and_acquire(self, key: str, value: Any) -> None:
+        # The factory can be long to complete (that's one of the reason we want to
+        # cache its results). Because we don't lock the factory call, together with
+        # the cache update, in case of concurrent access, quicker factories will
+        # complete first and will stay longer in the cache. This is fine, the
+        # alternative is to lock factory + cache update to avoid concurrent access but
+        # this will mean waiting for the first factory call to complete, this can
+        # potentially imply longer waits than no getting the value from the cache
+        with self._lock:
+            self._acquire(value)
+        # Can evict other resources, must be called without holding the lock
+        self._cache[key] = value
+
+    def _acquire(self, value: Any) -> None:
+        value_id = id(value)
+        self._n_leases[value_id] = self._n_leases.get(value_id, 0) + 1
+
+    def _release(self, value: Any) -> None:
+        value_id = id(value)
+        with self._lock:
+            self._n_leases[value_id] -= 1
+            if self._n_leases[value_id]:
+                return
+            del self._n_leases[value_id]
+            pending = self._pending_evictions.pop(value_id, None)
+        if pending is not None:
+            self._evict(*pending)
+
+    def _on_evict(self, key: Any, value: Any) -> None:
+        with self._lock:
+            if id(value) in self._n_leases:
+                self._pending_evictions[id(value)] = (key, value)
+                return
+        self._evict(key, value)
+
+    def _evict(self, key: Any, value: Any) -> None:
+        if self._eviction_callback is not None:
+            self._eviction_callback(key, value)
+
 
 def close_cm_callback(key: str, value: Any) -> None:  # noqa: ARG001
     if hasattr(value, "__exit__"):

{datashare_python-0.11.2.dist-info → datashare_python-0.11.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: datashare-python
-Version: 0.11.2
+Version: 0.11.4
 Summary: Manage Python tasks and local resources in Datashare
 Project-URL: Homepage, https://icij.github.io/datashare-python/
 Project-URL: Documentation, https://icij.github.io/datashare-python/

{datashare_python-0.11.2.dist-info → datashare_python-0.11.4.dist-info}/RECORD

@@ -14,15 +14,15 @@ datashare_python/objects.py,sha256=RuDEC4dZPndMYwFKhmMlToY8jjwjupG7mGGCG5tzc6g,1
 datashare_python/task_client.py,sha256=oTmP8bvZW0UyhLNMi1AV3XIAx7hrdbxNRss2Mw2azEc,8435
 datashare_python/template.py,sha256=N1uM8TRXh5tQPV7BC1oJj-Pdm2Gs_SkjaXv50RnNYso,4682
 datashare_python/types_.py,sha256=cPcCkIk6ecILrKetPnAsrb0njHYuWJU3eZicUXy2lF8,1278
-datashare_python/utils.py,sha256=k85QP34hSHBRFo8YwGA7rYPO5cqCt33i3ntl2oSfHkE,28603
-datashare_python/worker-template.tar.gz,sha256=3Lu6zLEXlndprxXPGMN9L19XrplmxUki3v15HymMmls,236583
+datashare_python/utils.py,sha256=MU_laix9TclXQ7Pt0mqqYQylYR3OKU-MItjQupcG3c8,30183
+datashare_python/worker-template.tar.gz,sha256=5U1vAMg7xtJSKAz20kqz7QLknvzmREN5ddqK7RvAYyI,236093
 datashare_python/worker.py,sha256=ad36_yJOgwtj8KBheWndGkqaP_2aXTPSrYaHrbCXTvg,8389
 datashare_python/cli/__init__.py,sha256=9BPWtssDgsVfWMsZ1TtZCla0EC_kai4RHttr8oNLYOE,1401
 datashare_python/cli/project.py,sha256=w32Gy9AOL5B00uDT4in7YUCt2g68FnNbvwg2M3a8G6o,946
 datashare_python/cli/task.py,sha256=GsyR5o0fN4Pd8tJKAUgivAvCBKT3L4Ay5j_BCW1XBHk,5828
 datashare_python/cli/utils.py,sha256=p69CQb0zfixuyBkiZprhdMCc_NuYwXyAn6vC9H1UzAw,911
 datashare_python/cli/worker.py,sha256=x0p3a1_NsZPbmV7-WSVMRyW3_iadYcgyZKbDvWgymQc,5790
-datashare_python-0.11.2.dist-info/METADATA,sha256=STmFFjil2tcX6BgZOyEBdlraEtIs69qEWV6FGHGhGnY,1012
-datashare_python-0.11.2.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
-datashare_python-0.11.2.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
-datashare_python-0.11.2.dist-info/RECORD,,
+datashare_python-0.11.4.dist-info/METADATA,sha256=wesn34A1Y6Omeo3gftg-Jd4jiXo4a4hktN1blM3vCJI,1012
+datashare_python-0.11.4.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+datashare_python-0.11.4.dist-info/entry_points.txt,sha256=ILE7auxabHWiu3GC-AunWnzjhOI_SbZp7D4GqZHlLw4,68
+datashare_python-0.11.4.dist-info/RECORD,,

{datashare_python-0.11.2.dist-info → datashare_python-0.11.4.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.32.0
+Generator: hatchling 1.32.4
 Root-Is-Purelib: true
 Tag: py3-none-any