From 55bf72aead4fd27f49ecec4b4ef0cfbfc74733cc Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Tue, 20 Apr 2021 15:55:55 +0300
Subject: [PATCH] chore(prophet): bump prophet to 1.0.1 (#14228)

---
 setup.py                                |  2 +-
 superset/utils/pandas_postprocessing.py |  6 +++---
 tests/charts/api_tests.py               |  2 +-
 tests/pandas_postprocessing_tests.py    | 11 ++++++++++-
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index 68e33df6c8..ff01c686ac 100644
--- a/setup.py
+++ b/setup.py
@@ -145,7 +145,7 @@ setup(
         "postgres": ["psycopg2-binary==2.8.5"],
         "presto": ["pyhive[presto]>=0.4.0"],
         "trino": ["sqlalchemy-trino>=0.2"],
-        "prophet": ["fbprophet>=0.7.1, <0.8", "pystan<3.0"],
+        "prophet": ["prophet>=1.0.1, <1.1", "pystan<3.0"],
         "redshift": ["sqlalchemy-redshift>=0.8.1, < 0.9"],
         "snowflake": ["snowflake-sqlalchemy>=1.2.3, <1.3"],
         "teradata": ["sqlalchemy-teradata==0.9.0.dev0"],
diff --git a/superset/utils/pandas_postprocessing.py b/superset/utils/pandas_postprocessing.py
index 6a5a74e697..f352e5fad4 100644
--- a/superset/utils/pandas_postprocessing.py
+++ b/superset/utils/pandas_postprocessing.py
@@ -631,14 +631,14 @@ def _prophet_fit_and_predict(  # pylint: disable=too-many-arguments
     Fit a prophet model and return a DataFrame with predicted results.
     """
     try:
-        prophet_logger = logging.getLogger("fbprophet.plot")
+        prophet_logger = logging.getLogger("prophet.plot")
 
         prophet_logger.setLevel(logging.CRITICAL)
-        from fbprophet import Prophet  # pylint: disable=import-error
+        from prophet import Prophet  # pylint: disable=import-error
 
         prophet_logger.setLevel(logging.NOTSET)
     except ModuleNotFoundError:
-        raise QueryObjectValidationError(_("`fbprophet` package not installed"))
+        raise QueryObjectValidationError(_("`prophet` package not installed"))
     model = Prophet(
         interval_width=confidence_interval,
         yearly_seasonality=yearly_seasonality,
diff --git a/tests/charts/api_tests.py b/tests/charts/api_tests.py
index 6c9cfe1a5e..d3b1c9753d 100644
--- a/tests/charts/api_tests.py
+++ b/tests/charts/api_tests.py
@@ -1241,7 +1241,7 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
         """
         Chart data API: Ensure prophet post transformation works
         """
-        pytest.importorskip("fbprophet")
+        pytest.importorskip("prophet")
         self.login(username="admin")
         request_payload = get_query_context("birth_names")
         time_grain = "P1Y"
diff --git a/tests/pandas_postprocessing_tests.py b/tests/pandas_postprocessing_tests.py
index 030b177526..73d3d7c94b 100644
--- a/tests/pandas_postprocessing_tests.py
+++ b/tests/pandas_postprocessing_tests.py
@@ -16,6 +16,7 @@
 # under the License.
 # isort:skip_file
 from datetime import datetime
+from importlib.util import find_spec
 import math
 from typing import Any, List, Optional
 
@@ -560,7 +561,7 @@ class TestPostProcessing(SupersetTestCase):
         self.assertListEqual(processed_df["pct_a"].tolist(), [0.25, 0.75])
 
     def test_prophet_valid(self):
-        pytest.importorskip("fbprophet")
+        pytest.importorskip("prophet")
 
         df = proc.prophet(
             df=prophet_df, time_grain="P1M", periods=3, confidence_interval=0.9
@@ -588,6 +589,14 @@ class TestPostProcessing(SupersetTestCase):
         assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 5, 31)
         assert len(df) == 9
 
+    def test_prophet_import(self):
+        prophet = find_spec("prophet")
+        if prophet is None:
+            with pytest.raises(QueryObjectValidationError):
+                proc.prophet(
+                    df=prophet_df, time_grain="P1M", periods=3, confidence_interval=0.9
+                )
+
     def test_prophet_missing_temporal_column(self):
         df = prophet_df.drop(DTTM_ALIAS, axis=1)