Merge pull request #12 from neuro-ml/dev

`fetch` is a generator now

.github/workflows/tests.yml

@@ -7,7 +7,7 @@ env:
 
 jobs:
   test:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-20.04
     strategy:
       matrix:
         python-version: [ 3.6, 3.7, 3.8, 3.9, '3.10' ]

tarn/__version__.py

@@ -1 +1 @@
-__version__ = '0.1.3'
+__version__ = '0.2.0'

tarn/interface.py

@@ -69,7 +69,7 @@ class RemoteStorage:
 
     @abstractmethod
     def fetch(self, keys: Sequence[Key], store: Callable[[Key, Path], Any],
-              config: HashConfig) -> Sequence[Tuple[Any, bool]]:
+              config: HashConfig) -> Iterable[Tuple[Any, bool]]:
         """
         Fetches the value for ``key`` from a remote location.
         """
@@ -143,34 +143,35 @@ def read(self, key, context, *, fetch: bool) -> Tuple[Any, Union[None, bool]]:
 
         return None, status
 
-    def fetch(self, keys: Iterable[Key], context, *, verbose: bool) -> Sequence[Key]:
+    def fetch(self, keys: Sequence[Key], context, *, verbose: bool) -> Iterable[Key]:
+        """ Fetch the `keys` from remote. Yields the keys that were successfully fetched """
+
         def store(k, base):
             status = self._replicate(k, base, context, self.levels)
             bar.update()
             return status if status is WriteError else k
 
         keys = set(keys)
         bar = tqdm(disable=not verbose, total=len(keys))
-        present = set()
-        for key in keys:
+        present = 0
+        for key in list(keys):
             if self._contains(key, context):
-                present.add(key)
+                present += 1
+                keys.remove(key)
                 bar.update()
+                yield key
 
-        keys -= present
-        logger.info(f'Fetch: {len(present)} keys already present, fetching {len(keys)}')
+        logger.info('%s keys already present, fetching %s', present, len(keys))
 
         for remote in self.remote:
             if not keys:
                 break
 
-            logger.info(f'Trying remote {remote}')
-            keys -= {
-                k for k, success in remote.fetch(list(keys), store, self.hash)
-                if success and k is not WriteError
-            }
-
-        return list(keys)
+            logger.info('Trying remote %s', remote)
+            for key, success in remote.fetch(list(keys), store, self.hash):
+                if success and key is not WriteError:
+                    keys.remove(key)
+                    yield key
 
     def _contains(self, key, context):
         for layer in self.levels:

tarn/local/storage.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 from pathlib import Path
 from typing import Sequence, Iterable, Callable, Union
 
@@ -50,8 +51,22 @@ def write(self, file: PathLike, error: bool = True) -> Union[Key, None]:
 
         return key
 
-    def fetch(self, keys: Iterable[Key], *, verbose: bool) -> Sequence[Key]:
-        return self.storage.fetch(keys, None, verbose=verbose)
+    def fetch(self, keys: Sequence[Key], *, verbose: bool, legacy: bool = True) -> Iterable[Key]:
+        """ Fetch the `keys` from remote. Yields the keys that were successfully fetched """
+        keys = list(keys)
+        result = self.storage.fetch(keys, None, verbose=verbose)
+        if legacy:
+            warnings.warn(
+                'In a future release fetch will yield the successfully processed keys. '
+                'Pass legacy=False to adopt this behaviour early on', UserWarning,
+            )
+            warnings.warn(
+                'In a future release fetch will yield the successfully processed keys. '
+                'Pass legacy=False to adopt this behaviour early on', DeprecationWarning,
+            )
+            result = list(set(keys) - set(result))
+
+        return result
 
     def resolve(self, key: Key, *, fetch: bool = True) -> Path:
         """ This is not safe, but it's fast. """

tarn/remote/http.py

@@ -1,7 +1,7 @@
 import shutil
 import tempfile
 from pathlib import Path
-from typing import Sequence, Callable, Any, Tuple
+from typing import Sequence, Callable, Any, Tuple, Iterable
 from urllib.error import HTTPError, URLError
 from urllib.parse import urljoin
 from urllib.request import urlretrieve
@@ -23,29 +23,27 @@ def __init__(self, url: str, optional: bool = False):
         self.levels = self.hash = None
 
     def fetch(self, keys: Sequence[Key], store: Callable[[Key, Path], Any],
-              config: HashConfig) -> Sequence[Tuple[Any, bool]]:
+              config: HashConfig) -> Iterable[Tuple[Any, bool]]:
 
-        results = []
         with tempfile.TemporaryDirectory() as temp_dir:
             source = Path(temp_dir) / 'source'
             if keys and not self._get_config(config):
-                return [(None, False)] * len(keys)
+                yield from [(None, False)] * len(keys)
+                return
 
             for key in keys:
                 try:
                     self._fetch_tree(key_to_relative(key, self.levels), source)
 
                     value = store(key, source)
                     shutil.rmtree(source)
-                    results.append((value, True))
+                    yield value, True
 
                 except requests.exceptions.ConnectionError:
-                    results.append((None, False))
+                    yield None, False
 
                 shutil.rmtree(source, ignore_errors=True)
 
-        return results
-
     def _fetch_one(self, relative, local):
         try:
             urlretrieve(urljoin(self.url, str(relative)), str(local))

tarn/remote/ssh.py

@@ -2,7 +2,7 @@
 import socket
 import tempfile
 from pathlib import Path
-from typing import Union, Sequence, Callable, Any, Tuple
+from typing import Union, Sequence, Callable, Any, Tuple, Iterable
 
 import paramiko
 from paramiko import SSHClient, AuthenticationException, SSHException
@@ -49,7 +49,7 @@ def __init__(self, hostname: str, root: PathLike, port: int = SSH_PORT, username
         self.optional = optional
 
     def fetch(self, keys: Sequence[Key], store: Callable[[Key, Path], Any],
-              config: HashConfig) -> Sequence[Tuple[Any, bool]]:
+              config: HashConfig) -> Iterable[Tuple[Any, bool]]:
 
         try:
             self.ssh.connect(
@@ -64,36 +64,35 @@ def fetch(self, keys: Sequence[Key], store: Callable[[Key, Path], Any],
             if not self.optional:
                 raise
 
-            return [(None, False)] * len(keys)
+            yield from [(None, False)] * len(keys)
+            return
 
-        results = []
         try:
             with SCPClient(self.ssh.get_transport()) as scp, tempfile.TemporaryDirectory() as temp_dir:
                 source = Path(temp_dir) / 'source'
                 if keys and not self._get_config(scp, config):
-                    return [(None, False)] * len(keys)
+                    yield from [(None, False)] * len(keys)
+                    return
 
                 for key in keys:
                     try:
                         scp.get(str(self.root / key_to_relative(key, self.levels)), str(source), recursive=True)
                         if not source.exists():
-                            results.append((None, False))
+                            yield None, False
 
                         else:
                             value = store(key, source)
                             shutil.rmtree(source)
-                            results.append((value, True))
+                            yield value, True
 
                     except (SCPException, socket.timeout, SSHException):
-                        results.append((None, False))
+                        yield None, False
 
                     shutil.rmtree(source, ignore_errors=True)
 
         finally:
             self.ssh.close()
 
-        return results
-
     def _get_config(self, scp, config):
         try:
             if self.levels is None:

tests/test_remote/test_http.py

@@ -38,12 +38,14 @@ def test_missing(storage_factory):
 
         key = remote.write(__file__)
         missing = key[:-1] + 'x'
-        result = location.fetch([key, missing], lambda k, base: (base / 'data').exists(), location.hash)
+        result = list(location.fetch([key, missing], lambda k, base: (base / 'data').exists(), location.hash))
         assert result == [(True, True), (None, False)]
 
 
 def test_wrong_address():
     with pytest.raises(requests.exceptions.ConnectionError):
-        HTTPLocation('http://localhost/wrong').fetch(['some-key'], lambda *args: True, None)
+        list(HTTPLocation('http://localhost/wrong').fetch(['some-key'], lambda *args: True, None))
 
-    assert HTTPLocation('http://localhost/wrong', True).fetch(['some-key'], lambda *args: True, None) == [(None, False)]
+    assert list(HTTPLocation(
+        'http://localhost/wrong', True
+    ).fetch(['some-key'], lambda *args: True, None)) == [(None, False)]

tests/test_remote/test_ssh.py

@@ -31,6 +31,7 @@ def test_storage_ssh(storage_factory):
 
 def test_wrong_host():
     with pytest.raises((NoValidConnectionsError, SSHException)):
-        SSHLocation('localhost', '/').fetch(['some-key'], lambda *args: True, None)
-    assert SSHLocation('localhost', '/', optional=True).fetch(
-        ['some-key'], lambda *args: True, None) == [(None, False)]
+        list(SSHLocation('localhost', '/').fetch(['some-key'], lambda *args: True, None))
+    assert list(SSHLocation(
+        'localhost', '/', optional=True
+    ).fetch(['some-key'], lambda *args: True, None)) == [(None, False)]