(superset) 01/01: chore: remove prophet error

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch supress_prophet_logging
in repository https://gitbox.apache.org/repos/asf/superset.git

commit f4dca0e55998769dc5365ce792f37131471f113b
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Thu Feb 8 12:56:36 2024 -0500

    chore: remove prophet error
---
 superset/utils/decorators.py                    | 19 +++++++++++++++++++
 superset/utils/pandas_postprocessing/prophet.py |  7 +++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py
index 43b6909caf..df66b707bb 100644
--- a/superset/utils/decorators.py
+++ b/superset/utils/decorators.py
@@ -191,3 +191,22 @@ def debounce(duration: float | int = 0.1) -> Callable[..., Any]:
 
 def on_security_exception(self: Any, ex: Exception) -> Response:
     return self.response(403, **{"message": utils.error_msg_from_exception(ex)})
+
+
+@contextmanager
+def suppress_logging(
+    logger_name: str | None = None,
+    level: int = logging.CRITICAL,
+) -> Iterator[None]:
+    """
+    Context manager to suppress logging during the execution of code block.
+
+    Use with caution and make sure you have the least amount of code inside it.
+    """
+    target_logger = logging.getLogger(logger_name)
+    original_level = target_logger.getEffectiveLevel()
+    target_logger.setLevel(level)
+    try:
+        yield
+    finally:
+        target_logger.setLevel(original_level)
diff --git a/superset/utils/pandas_postprocessing/prophet.py b/superset/utils/pandas_postprocessing/prophet.py
index 47d956fed5..3d7fb42b8b 100644
--- a/superset/utils/pandas_postprocessing/prophet.py
+++ b/superset/utils/pandas_postprocessing/prophet.py
@@ -23,6 +23,7 @@ from pandas import DataFrame
 
 from superset.exceptions import InvalidPostProcessingError
 from superset.utils.core import DTTM_ALIAS
+from superset.utils.decorators import suppress_logging
 from superset.utils.pandas_postprocessing.utils import PROPHET_TIME_GRAIN_MAP
 
 
@@ -52,8 +53,10 @@ def _prophet_fit_and_predict(  # pylint: disable=too-many-arguments
     Fit a prophet model and return a DataFrame with predicted results.
     """
     try:
-        # pylint: disable=import-outside-toplevel
-        from prophet import Prophet
+        # `prophet` complains about `plotly` not being installed
+        with suppress_logging("prophet.plot", logging.CRITICAL):
+            # pylint: disable=import-outside-toplevel
+            from prophet import Prophet
 
         prophet_logger = logging.getLogger("prophet.plot")
         prophet_logger.setLevel(logging.CRITICAL)