Do not display the download bar if we are streaming URL's

Author: gordonwatts

scripts/test.ipynb

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from datetime import timedelta
 from pathlib import Path
-from typing import (Any, AsyncIterator, Awaitable, Callable, Dict, List,
+from typing import (Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, Dict, List,
                     Optional, Tuple, Union)
 
 import aiohttp
@@ -287,7 +287,7 @@ async def convert_to_file(f: Path) -> Path:
 
     @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
     async def _stream_url_buckets(self, selection_query: str, data_format: str) \
-            -> AsyncIterator[StreamInfoUrl]:
+            -> AsyncGenerator[StreamInfoUrl, None]:
         '''Get a list of files back for a request
 
         Args:
@@ -314,7 +314,7 @@ async def _stream_url_buckets(self, selection_query: str, data_format: str) \
             # Look up the cache, and then fetch an iterator going thorugh the results
             # from either servicex or the cache, depending.
             try:
-                notifier = self._create_notifier()
+                notifier = self._create_notifier(False)
                 minio_files = self._get_minio_bucket_files_from_servicex(request_id, client,
                                                                          minio_adaptor, notifier)
 
@@ -405,7 +405,8 @@ async def _stream_return(self, selection_query: str,
 
     @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
     async def _stream_local_files(self, selection_query: str,
-                                  data_format: str = 'root-file'):
+                                  data_format: str = 'root-file') \
+            -> AsyncGenerator[StreamInfoPath, None]:
         '''
         Given a query, return the data as a list of paths pointing to local files
         that contain the results of the query. This is an async generator, and files
@@ -425,7 +426,7 @@ async def _stream_local_files(self, selection_query: str,
                                 on the converter call.
         '''
         # Get a notifier to update anyone who wants to listen.
-        notifier = self._create_notifier()
+        notifier = self._create_notifier(True)
 
         # Get all the files
         as_files = \

servicex/servicex.py

@@ -68,9 +68,9 @@ def __init__(self,
             status_callback_factory if status_callback_factory is not None \
             else _null_progress_feedback
 
-    def _create_notifier(self) -> _status_update_wrapper:
+    def _create_notifier(self, downloading: bool) -> _status_update_wrapper:
         'Internal method to create a updater from the status call-back'
-        return _status_update_wrapper(self._status_callback_factory(self._dataset))
+        return _status_update_wrapper(self._status_callback_factory(self._dataset, downloading))
 
     @abstractmethod
     async def get_data_rootfiles_async(self, selection_query: str) -> List[Path]:

servicex/servicexabc.py

@@ -87,11 +87,16 @@ def sanitize_filename(fname: str):
 
 
 StatusUpdateCallback = Callable[[Optional[int], int, int, int], None]
-# The sig of the call-back
+# The sig of the call-back. Arguments are:
+# 1. Processed
+# 1. Downloaded
+# 1. Remaining
+# 1. Failed
 
-StatusUpdateFactory = Callable[[str], StatusUpdateCallback]
+StatusUpdateFactory = Callable[[str, bool], StatusUpdateCallback]
 # Factory method that returns a factory that can run the status callback
-# First argument is a string.
+# First argument is a string, second argument is True if a download progress
+# bar can be displayed as well.
 
 
 class _status_update_wrapper:
@@ -192,14 +197,20 @@ async def stream_unique_updates_only(stream: AsyncIterator[TransformTuple]):
             yield p
 
 
-def _run_default_wrapper(ds_name: str) -> StatusUpdateCallback:
-    '''
-    Create a feedback object for everyone to use to pass feedback to. Uses tqdm (default).
+def _run_default_wrapper(ds_name: str, downloading: bool) -> StatusUpdateCallback:
+    '''Create a feedback object for everyone to use to pass feedback to. Uses tqdm (default).
+
+    Args:
+        ds_name (str): The dataset name
+        downloading (bool): True if we will be downloading the files, not just processing them.
+
+    Returns:
+        StatusUpdateCallback: The updater callback
     '''
-    return _default_wrapper_mgr(ds_name).update
+    return _default_wrapper_mgr(ds_name, downloading).update
 
 
-def _null_progress_feedback(ds_name: str) -> None:
+def _null_progress_feedback(ds_name: str, downloading: bool) -> None:
     '''
     Internal routine to create a feedback object that does not
     give anyone feedback!
@@ -209,9 +220,10 @@ def _null_progress_feedback(ds_name: str) -> None:
 
 class _default_wrapper_mgr:
     'Default progress bar'
-    def __init__(self, sample_name: Optional[str] = None):
+    def __init__(self, sample_name: Optional[str] = None, show_download_bar: bool = True):
         self._tqdm_p: Optional[tqdm] = None
         self._tqdm_d: Optional[tqdm] = None
+        self._show_download_progress = show_download_bar
         self._sample_name = sample_name
 
     def _init_tqdm(self):
@@ -221,9 +233,10 @@ def _init_tqdm(self):
         self._tqdm_p = tqdm(total=9e9, desc=self._sample_name, unit='file',
                             leave=False, dynamic_ncols=True,
                             bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]')
-        self._tqdm_d = tqdm(total=9e9, desc="        Downloaded", unit='file',
-                            leave=False, dynamic_ncols=True,
-                            bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]')
+        if self._show_download_progress:
+            self._tqdm_d = tqdm(total=9e9, desc="        Downloaded", unit='file',
+                                leave=False, dynamic_ncols=True,
+                                bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]')
 
     def _update_bar(self, bar: Optional[tqdm], total: Optional[int], num: int, failed: int):
         assert bar is not None, 'Internal error - bar was not initalized'
@@ -244,7 +257,8 @@ def _update_bar(self, bar: Optional[tqdm], total: Optional[int], num: int, faile
     def update(self, total: Optional[int], processed: int, downloaded: int, failed: int):
         self._init_tqdm()
         self._update_bar(self._tqdm_p, total, processed, failed)
-        self._update_bar(self._tqdm_d, total, downloaded, failed)
+        if self._show_download_progress:
+            self._update_bar(self._tqdm_d, total, downloaded, failed)
 
 
 # Changes in the json that won't affect the result

servicex/utils.py

@@ -778,7 +778,7 @@ def check_in(total: Optional[int], processed: int, downloaded: int, failed: int)
                             cache_adaptor=mock_cache,
                             data_convert_adaptor=data_adaptor,
                             local_log=mock_logger,
-                            status_callback_factory=lambda ds: check_in)
+                            status_callback_factory=lambda ds, downloading: check_in)
     ds.get_data_rootfiles('(valid qastle string)')
 
     assert f_total == 1
@@ -787,6 +787,73 @@ def check_in(total: Optional[int], processed: int, downloaded: int, failed: int)
     assert f_failed == 0
 
 
+def test_callback_is_downloading(mocker):
+    'Make sure this file download sets the file-download marker in the callback'
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+
+    def check_in(total: Optional[int], processed: int, downloaded: int, failed: int):
+        pass
+
+    ds_name = None
+    ds_downloading = None
+
+    def build_it(ds: str, downloading: bool):
+        nonlocal ds_name, ds_downloading
+        ds_name = ds
+        ds_downloading = downloading
+        return check_in
+
+    ds = fe.ServiceXDataset('http://one-ds',
+                            servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+                            minio_adaptor=mock_minio_adaptor,  # type: ignore
+                            cache_adaptor=mock_cache,
+                            data_convert_adaptor=data_adaptor,
+                            local_log=mock_logger,
+                            status_callback_factory=build_it)
+    ds.get_data_rootfiles('(valid qastle string)')
+
+    assert ds_name == 'http://one-ds'
+    assert ds_downloading
+
+
+@pytest.mark.asyncio
+async def test_callback_is_not_downloading(mocker):
+    'Stream download of Minio URLs should not set the download marker'
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+
+    def check_in(total: Optional[int], processed: int, downloaded: int, failed: int):
+        pass
+
+    ds_name = None
+    ds_downloading = None
+
+    def build_it(ds: str, downloading: bool):
+        nonlocal ds_name, ds_downloading
+        ds_name = ds
+        ds_downloading = downloading
+        return check_in
+
+    ds = fe.ServiceXDataset('http://one-ds',
+                            servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+                            minio_adaptor=mock_minio_adaptor,  # type: ignore
+                            cache_adaptor=mock_cache,
+                            data_convert_adaptor=data_adaptor,
+                            local_log=mock_logger,
+                            status_callback_factory=build_it)
+    _ = [f async for f in ds.get_data_rootfiles_url_stream('(valid qastle string)')]
+
+    assert ds_name == 'http://one-ds'
+    assert not ds_downloading
+
+
 @pytest.mark.asyncio
 async def test_callback_none(mocker):
     'Get a root file with a single file'