Move the fancy rich prints over to being logging calls.

Author: gordonwatts

servicex/query.py

@@ -31,6 +31,7 @@
 import asyncio
 from abc import ABC
 from asyncio import Task, CancelledError
+import logging
 from typing import List, Optional, Union
 from servicex.expandable_progress import ExpandableProgress
 
@@ -43,9 +44,7 @@
 from servicex.types import DID
 
 import rich
-from rich.progress import (
-    Progress,
-    TaskID)
+from rich.progress import Progress, TaskID
 
 from servicex.configuration import Configuration
 from servicex.minio_adapter import MinioAdapter
@@ -157,9 +156,10 @@ def set_result_format(self, result_format: ResultFormat):
         return self
 
     async def submit_and_download(
-        self, signed_urls_only: bool,
-            expandable_progress: ExpandableProgress,
-            dataset_group: Optional[bool] = False,
+        self,
+        signed_urls_only: bool,
+        expandable_progress: ExpandableProgress,
+        dataset_group: Optional[bool] = False,
     ) -> Optional[TransformedResults]:
         """
         Submit the transform request to ServiceX. Poll the transform status to see when
@@ -186,27 +186,31 @@ def transform_complete(task: Task):
             :return:
             """
             if task.exception():
-                rich.print("ServiceX Exception", task.exception())
+                logging.getLogger(__name__).error(
+                    "ServiceX Exception", exc_info=task.exception()
+                )
                 if download_files_task:
                     download_files_task.cancel("Transform failed")
                 raise task.exception()
 
             if self.current_status.files_failed:
-                rich.print(
-                    f"[bold red]Transforms completed with failures[/bold red] "
+                logging.getLogger(__name__).warning(
+                    f"Transforms completed with failures "
                     f"{self.current_status.files_failed} files failed out of "
                     f"{self.current_status.files}"
                 )
             else:
-                rich.print("Transforms completed successfully")
+                logging.getLogger(__name__).info("Transforms completed successfully")
 
         sx_request = self.transform_request
 
         # Let's see if this is in the cache already, but respect the user's wishes
         # to ignore the cache
-        cached_record = self.cache.get_transform_by_hash(sx_request.compute_hash()) \
-            if not self.ignore_cache \
+        cached_record = (
+            self.cache.get_transform_by_hash(sx_request.compute_hash())
+            if not self.ignore_cache
             else None
+        )
 
         # And that we grabbed the resulting files in the way that the user requested
         # (Downloaded, or obtained pre-signed URLs)
@@ -217,17 +221,21 @@ def transform_complete(task: Task):
                 or not signed_urls_only
                 and cached_record.file_list
             ):
-                rich.print("Returning results from cache")
+                logging.getLogger(__name__).info("Returning results from cache")
                 return cached_record
 
         # If we get here with a cached record, then we know that the transform
         # has been run, but we just didn't get the files from object store in the way
         # requested by user
         transform_bar_title = f"{sx_request.title}: Transform"
         if not cached_record:
-            transform_progress = expandable_progress.add_task(
-                transform_bar_title, start=False, total=None
-            ) if expandable_progress else None
+            transform_progress = (
+                expandable_progress.add_task(
+                    transform_bar_title, start=False, total=None
+                )
+                if expandable_progress
+                else None
+            )
         else:
             self.request_id = cached_record.request_id
             transform_progress = None
@@ -236,19 +244,28 @@ def transform_complete(task: Task):
         minio_progress_bar_title = (
             "Download" if not signed_urls_only else "Signing URLS"
         )
-        minio_progress_bar_title = minio_progress_bar_title.rjust(len(transform_bar_title))
+        minio_progress_bar_title = minio_progress_bar_title.rjust(
+            len(transform_bar_title)
+        )
 
-        download_progress = expandable_progress.add_task(
-            minio_progress_bar_title, start=False, total=None
-        ) if expandable_progress else None
+        download_progress = (
+            expandable_progress.add_task(
+                minio_progress_bar_title, start=False, total=None
+            )
+            if expandable_progress
+            else None
+        )
 
         if not cached_record:
             self.request_id = await self.servicex.submit_transform(sx_request)
 
             monitor_task = loop.create_task(
                 self.transform_status_listener(
-                    expandable_progress, transform_progress, transform_bar_title,
-                    download_progress, minio_progress_bar_title
+                    expandable_progress,
+                    transform_progress,
+                    transform_bar_title,
+                    download_progress,
+                    minio_progress_bar_title,
                 )
             )
             monitor_task.add_done_callback(transform_complete)
@@ -289,11 +306,17 @@ def transform_complete(task: Task):
 
             return transform_report
         except CancelledError:
-            rich.print("Aborted file downloads due to transform failure")
+            logging.getLogger(__name__).warning(
+                "Aborted file downloads due to transform failure"
+            )
 
     async def transform_status_listener(
-        self, progress: ExpandableProgress, progress_task: TaskID,
-        progress_bar_title: str, download_task: TaskID, download_bar_title: str
+        self,
+        progress: ExpandableProgress,
+        progress_task: TaskID,
+        progress_bar_title: str,
+        download_task: TaskID,
+        download_bar_title: str,
     ):
         """
         Poll ServiceX for the status of a transform. Update progress bars and keep track
@@ -314,16 +337,21 @@ async def transform_status_listener(
             if not final_count and self.current_status.files:
                 final_count = self.current_status.files
                 if progress:
-                    progress.update(progress_task, progress_bar_title, total=final_count)
+                    progress.update(
+                        progress_task, progress_bar_title, total=final_count
+                    )
                     progress.start_task(task_id=progress_task, task_type="Transform")
 
-                    progress.update(download_task, download_bar_title, total=final_count)
+                    progress.update(
+                        download_task, download_bar_title, total=final_count
+                    )
                     progress.start_task(task_id=download_task, task_type="Download")
 
             if progress:
                 progress.update(
-                    progress_task, progress_bar_title,
-                    completed=self.current_status.files_completed
+                    progress_task,
+                    progress_bar_title,
+                    completed=self.current_status.files_completed,
                 )
 
             if self.current_status.status == Status.complete:
@@ -339,7 +367,9 @@ async def retrieve_current_transform_status(self):
         # Is this the first time we've polled status? We now know the request ID.
         # Update the display and set our download directory.
         if not self.current_status:
-            rich.print(f"[bold]ServiceX Transform {s.title}: {s.request_id}[/bold]")
+            logging.getLogger(__name__).info(
+                f"ServiceX Transform {s.title}: {s.request_id}"
+            )
             self.download_path = self.cache.cache_path_for_transform(s)
 
         self.current_status = s
@@ -434,35 +464,39 @@ async def get_signed_url(
         await asyncio.gather(*download_tasks)
         return result_uris
 
-    async def as_files_async(self,
-                             display_progress: bool = True,
-                             provided_progress: Optional[ProgressIndicators] = None
-                             ) -> TransformedResults:
+    async def as_files_async(
+        self,
+        display_progress: bool = True,
+        provided_progress: Optional[ProgressIndicators] = None,
+    ) -> TransformedResults:
         r"""
         Submit the transform and request all the resulting files to be downloaded
         :return: TransformResult instance with the list of complete paths to the downloaded files
         """
         with ExpandableProgress(display_progress, provided_progress) as progress:
-            return await self.submit_and_download(signed_urls_only=False,
-                                                  expandable_progress=progress)
+            return await self.submit_and_download(
+                signed_urls_only=False, expandable_progress=progress
+            )
 
     as_files = make_sync(as_files_async)
 
     try:
-        async def as_pandas_async(self,
-                                  display_progress: bool = True,
-                                  provided_progress: Optional[ProgressIndicators] = None) \
-                -> pd.DataFrame:
+
+        async def as_pandas_async(
+            self,
+            display_progress: bool = True,
+            provided_progress: Optional[ProgressIndicators] = None,
+        ) -> pd.DataFrame:
             r"""
             Return a pandas dataframe containing the results. This only works if you've
             installed pandas extra
 
             :return: Pandas Dataframe
             """
             self.result_format = ResultFormat.parquet
-            transformed_result = await self.as_files_async(display_progress=display_progress,
-                                                           provided_progress=provided_progress
-                                                           )
+            transformed_result = await self.as_files_async(
+                display_progress=display_progress, provided_progress=provided_progress
+            )
             dataframes = pd.concat(
                 [pd.read_parquet(p) for p in transformed_result.file_list]
             )
@@ -472,24 +506,31 @@ async def as_pandas_async(self,
     except NameError:
         pass
 
-    async def as_signed_urls_async(self, display_progress: bool = True,
-                                   provided_progress: Optional[ProgressIndicators] = None,
-                                   dataset_group: bool = False) \
-            -> TransformedResults:
+    async def as_signed_urls_async(
+        self,
+        display_progress: bool = True,
+        provided_progress: Optional[ProgressIndicators] = None,
+        dataset_group: bool = False,
+    ) -> TransformedResults:
         r"""
         Presign URLs for each of the transformed files
 
         :return: TransformedResults object with the presigned_urls list populated
         """
         if dataset_group:
-            return await self.submit_and_download(signed_urls_only=True,
-                                                  expandable_progress=provided_progress,
-                                                  dataset_group=dataset_group)
-
-        with ExpandableProgress(display_progress=display_progress,
-                                provided_progress=provided_progress) as progress:
-            return await self.submit_and_download(signed_urls_only=True,
-                                                  expandable_progress=progress,
-                                                  dataset_group=dataset_group)
+            return await self.submit_and_download(
+                signed_urls_only=True,
+                expandable_progress=provided_progress,
+                dataset_group=dataset_group,
+            )
+
+        with ExpandableProgress(
+            display_progress=display_progress, provided_progress=provided_progress
+        ) as progress:
+            return await self.submit_and_download(
+                signed_urls_only=True,
+                expandable_progress=progress,
+                dataset_group=dataset_group,
+            )
 
     as_signed_urls = make_sync(as_signed_urls_async)

servicex/servicex_client.py

@@ -58,14 +58,19 @@ def get_codegen(_sample: Sample, _general: General):
     for sample in config.Sample:
         if sample.Query:
             if type(sample.Query) is str:
-                qastle_query = qastle.python_ast_to_text_ast(ast.parse(sample.Query)) # NOQA E501
+                qastle_query = qastle.python_ast_to_text_ast(
+                    ast.parse(sample.Query)
+                )  # NOQA E501
                 sample.Query = FuncADLQuery()
 
                 sample.Query.set_provided_qastle(qastle_query)
 
-            query = sx.func_adl_dataset(sample.dataset_identifier, sample.Name,
-                                        get_codegen(sample, config.General),
-                                        config.General.OutputFormat)
+            query = sx.func_adl_dataset(
+                sample.dataset_identifier,
+                sample.Name,
+                get_codegen(sample, config.General),
+                config.General.OutputFormat,
+            )
             query._q_ast = sample.Query._q_ast
             query._item_type = sample.Query._item_type
             if sample.Tree:
@@ -86,9 +91,12 @@ def get_codegen(_sample: Sample, _general: General):
                 except SyntaxError as e:
                     raise SyntaxError(f"Syntax error in {sample.Name}: {e}")
 
-            dataset = sx.python_dataset(sample.dataset_identifier, sample.Name,
-                                        get_codegen(sample, config.General),
-                                        config.General.OutputFormat)
+            dataset = sx.python_dataset(
+                sample.dataset_identifier,
+                sample.Name,
+                get_codegen(sample, config.General),
+                config.General.OutputFormat,
+            )
             dataset.python_function = sample.Function
             dataset.ignore_cache = sample.IgnoreLocalCache
             datasets.append(dataset)
@@ -173,7 +181,7 @@ def get_code_generators(self, backend=None):
         if backend:
             cached_backends = self.query_cache.get_codegen_by_backend(backend)
         if cached_backends:
-            rich.print("Returning code generators from cache")
+            logging.getLogger(__name__).info("Returning code generators from cache")
             return cached_backends["codegens"]
         else:
             code_generators = self.servicex.get_code_generators()
@@ -187,7 +195,7 @@ def func_adl_dataset(
         codegen: str = "uproot",
         result_format: Optional[ResultFormat] = None,
         item_type: Type[T] = Any,
-        ignore_cache: bool = False
+        ignore_cache: bool = False,
     ) -> FuncADLQuery[T]:
         r"""
         Generate a dataset that can use func_adl query language
@@ -217,7 +225,7 @@ def func_adl_dataset(
             query_cache=self.query_cache,
             result_format=result_format,
             item_type=item_type,
-            ignore_cache=ignore_cache
+            ignore_cache=ignore_cache,
         )
 
     def python_dataset(
@@ -226,7 +234,7 @@ def python_dataset(
         title: str = "ServiceX Client",
         codegen: str = "uproot",
         result_format: Optional[ResultFormat] = None,
-        ignore_cache: bool = False
+        ignore_cache: bool = False,
     ) -> PythonQuery:
         r"""
         Generate a dataset that can use accept a python function for the  query
@@ -256,5 +264,5 @@ def python_dataset(
             config=self.config,
             query_cache=self.query_cache,
             result_format=result_format,
-            ignore_cache=ignore_cache
+            ignore_cache=ignore_cache,
         )