add remaining annotations

Author: jorenham

python/prophet/diagnostics.py

@@ -42,7 +42,7 @@ class _SupportsMap(Protocol):
         def map(self, __f: Callable[..., object], *its: Iterable[object]) -> Any: ...
 
 
-logger = logging.getLogger('prophet')
+logger: logging.Logger = logging.getLogger('prophet')
 
 
 def generate_cutoffs(

python/prophet/forecaster.py

@@ -29,7 +29,7 @@
 from prophet.models import StanBackendEnum, ModelInputData, ModelParams, TrendIndicator, IStanBackend
 from prophet.plot import (plot, plot_components)
 
-logger = logging.getLogger('prophet')
+logger: logging.Logger = logging.getLogger('prophet')
 logger.setLevel(logging.INFO)
 NANOSECONDS_TO_SECONDS = 1000 * 1000 * 1000
 
@@ -91,13 +91,14 @@ class Prophet:
     growth: Literal["linear", "logistic", "flat"]
     changepoints: pd.Series[pd.Timestamp] | None
     n_changepoints: int
+    specified_changepoints: bool
     changepoint_range: float
     yearly_seasonality: Literal["auto"] | int
     weekly_seasonality: Literal["auto"] | int
     daily_seasonality: Literal["auto"] | int
     holidays: pd.DataFrame | None
-    seasonality_mode: Literal["additive", "multiplicative"]
-    holidays_mode: Literal["additive", "multiplicative"]
+    seasonality_mode: _Mode
+    holidays_mode: _Mode
     seasonality_prior_scale: float
     holidays_prior_scale: float
     changepoint_prior_scale: float
@@ -112,8 +113,8 @@ class Prophet:
     logistic_floor: bool
     t_scale: pd.Timedelta | None
     changepoints_t: npt.NDArray[np.float64] | None
-    seasonalities: OrderedDict[str, dict]
-    extra_regressors: OrderedDict[str, dict]
+    seasonalities: OrderedDict[str, dict[str, Any]]
+    extra_regressors: OrderedDict[str, dict[str, Any]]
     country_holidays: str | None
     stan_fit: Any | None
     params: dict[str, Any]
@@ -135,7 +136,7 @@ def __init__(
         weekly_seasonality: Literal["auto"] | int = "auto",
         daily_seasonality: Literal["auto"] | int = "auto",
         holidays: pd.DataFrame | None = None,
-        seasonality_mode: Literal["additive", "multiplicative"] = "additive",
+        seasonality_mode: _Mode = "additive",
         seasonality_prior_scale: SupportsFloat = 10.0,
         holidays_prior_scale: SupportsFloat = 10.0,
         changepoint_prior_scale: SupportsFloat = 0.05,
@@ -144,7 +145,7 @@ def __init__(
         uncertainty_samples: int = 1000,
         stan_backend: str | None = None,
         scaling: Literal["absmax", "minmax"] = "absmax",
-        holidays_mode: Literal["additive", "multiplicative"] | None = None,
+        holidays_mode: _Mode | None = None,
     ) -> None:
         self.growth = growth
 
@@ -690,7 +691,7 @@ def add_regressor(
         prior_scale: float | None = None,
         standardize: Literal['auto'] | bool = 'auto',
         mode: _Mode | None = None,
-    ):
+    ) -> Self:
         """Add an additional regressor to be used for fitting and predicting.
 
         The dataframe passed to `fit` and `predict` will have a column with the
@@ -1021,7 +1022,13 @@ def add_group_component(
             components = pd.concat([components, new_comp], ignore_index=True)
         return components
 
-    def parse_seasonality_args(self, name: str, arg, auto_disable: bool, default_order: int) -> int:
+    def parse_seasonality_args(
+        self,
+        name: str,
+        arg: Literal['auto'] | int,
+        auto_disable: bool,
+        default_order: int,
+    ) -> int:
         """Get number of fourier components for built-in seasonalities.
 
         Parameters
@@ -1201,7 +1208,7 @@ def flat_growth_init(df: pd.DataFrame) -> tuple[float, float]:
         m = df['y_scaled'].mean()
         return k, m
 
-    def preprocess(self, df: pd.DataFrame, **kwargs) -> ModelInputData:
+    def preprocess(self, df: pd.DataFrame, **kwargs: Any) -> ModelInputData:
         """
         Reformats historical data, standardizes y and extra regressors, sets seasonalities and changepoints.
 
@@ -1985,7 +1992,7 @@ def predictive_samples(self, df: pd.DataFrame, vectorized: bool = True) -> dict[
         df = self.setup_dataframe(df.copy())
         return self.sample_posterior_predictive(df, vectorized)
 
-    def percentile(self, a, *args, **kwargs):
+    def percentile(self, a: npt.ArrayLike, *args: Any, **kwargs: Any) -> np.ndarray:
         """
         We rely on np.nanpercentile in the rare instances where there
         are a small number of bad samples with MCMC that contain NaNs.

python/prophet/models.py

@@ -19,7 +19,7 @@
 import pandas as pd
 
 import logging
-logger = logging.getLogger('prophet.models')
+logger: logging.Logger = logging.getLogger('prophet.models')
 
 if TYPE_CHECKING:
     from cmdstanpy import CmdStanMCMC, CmdStanMLE, CmdStanModel

python/prophet/plot.py

@@ -16,16 +16,24 @@
 from prophet.diagnostics import performance_metrics
 
 if TYPE_CHECKING:
-    from typing import Literal, Sequence, TypeVar
+    from typing import Literal, Sequence, TypeVar, type_check_only
+    from typing_extensions import TypedDict
+    
     from prophet.forecaster import Prophet
 
     import matplotlib.pyplot as plt
     import plotly.graph_objs as go
 
     _AxT = TypeVar('_AxT', bound=plt.Axes)
 
+    @type_check_only
+    class _PlotlyProps(TypedDict):
+        traces: list[go.Scatter]
+        xaxis: go.layout.XAxis
+        yaxis: go.layout.YAxis
+
 
-logger = logging.getLogger('prophet.plot')
+logger: logging.Logger = logging.getLogger('prophet.plot')
 
 
 try:
@@ -798,7 +806,12 @@ def plot_plotly(
 
 
 def plot_components_plotly(
-        m, fcst, uncertainty=True, plot_cap=True, figsize=(900, 200)):
+    m: Prophet,
+    fcst: pd.DataFrame,
+    uncertainty: bool = True,
+    plot_cap: bool = True,
+    figsize: tuple[int, int] = (900, 200),
+) -> go.Figure:
     """Plot the Prophet forecast components using Plotly.
     See plot_plotly() for Plotly setup instructions
 
@@ -860,7 +873,14 @@ def plot_components_plotly(
     return fig
 
 
-def plot_forecast_component_plotly(m, fcst, name, uncertainty=True, plot_cap=False, figsize=(900, 300)):
+def plot_forecast_component_plotly(
+    m: Prophet,
+    fcst: pd.DataFrame,
+    name: str,
+    uncertainty: bool = True,
+    plot_cap: bool = False,
+    figsize: tuple[int, int] = (900, 300)
+) -> go.Figure:
     """Plot a particular component of the forecast using Plotly.
     See plot_plotly() for Plotly setup instructions
 
@@ -891,7 +911,12 @@ def plot_forecast_component_plotly(m, fcst, name, uncertainty=True, plot_cap=Fal
     return fig
 
 
-def plot_seasonality_plotly(m, name, uncertainty=True, figsize=(900, 300)):
+def plot_seasonality_plotly(
+    m: Prophet,
+    name: str,
+    uncertainty: bool = True,
+    figsize: tuple[int, int] = (900, 300)
+) -> go.Figure:
     """Plot a custom seasonal component using Plotly.
     See plot_plotly() for Plotly setup instructions
 
@@ -919,7 +944,13 @@ def plot_seasonality_plotly(m, name, uncertainty=True, figsize=(900, 300)):
     return fig
 
 
-def get_forecast_component_plotly_props(m, fcst, name, uncertainty=True, plot_cap=False):
+def get_forecast_component_plotly_props(
+    m: Prophet,
+    fcst: pd.DataFrame,
+    name: str,
+    uncertainty: bool = True,
+    plot_cap: bool = False,
+) -> _PlotlyProps:
     """Prepares a dictionary for plotting the selected forecast component with Plotly
 
     Parameters
@@ -956,8 +987,10 @@ def get_forecast_component_plotly_props(m, fcst, name, uncertainty=True, plot_ca
         holiday_features.columns = holiday_features.columns.str.replace('+0', '', regex=False)
         text = pd.Series(data='', index=holiday_features.index)
         for holiday_feature, idxs in holiday_features.items():
+            # https://github.com/facebook/pyrefly/issues/2248
+            # pyrefly:ignore[unsupported-operation]
             text[idxs.astype(bool) & (text != '')] += '<br>'  # Add newline if additional holiday
-            text[idxs.astype(bool)] += holiday_feature
+            text[idxs.astype(bool)] += holiday_feature  # pyrefly:ignore[unsupported-operation]
 
     traces = []
     traces.append(go.Scatter(
@@ -1020,12 +1053,17 @@ def get_forecast_component_plotly_props(m, fcst, name, uncertainty=True, plot_ca
     yaxis = go.layout.YAxis(rangemode='normal' if name == 'trend' else 'tozero',
                             title=go.layout.yaxis.Title(text=name),
                             zerolinecolor=zeroline_color)
+    assert m.component_modes
     if name in m.component_modes['multiplicative']:
         yaxis.update(tickformat='%', hoverformat='.2%')
     return {'traces': traces, 'xaxis': xaxis, 'yaxis': yaxis}
 
 
-def get_seasonality_plotly_props(m, name, uncertainty=True):
+def get_seasonality_plotly_props(
+    m: Prophet,
+    name: str,
+    uncertainty: bool = True,
+) -> _PlotlyProps:
     """Prepares a dictionary for plotting the selected seasonality with Plotly
 
     Parameters
@@ -1048,6 +1086,7 @@ def get_seasonality_plotly_props(m, name, uncertainty=True):
     start = pd.to_datetime('2017-01-01 0000')
     period = m.seasonalities[name]['period']
     end = start + pd.Timedelta(days=period)
+    assert m.history is not None
     if (m.history['ds'].dt.hour == 0).all():  # Day Precision
         plot_points = np.floor(period).astype(int)
     elif (m.history['ds'].dt.minute == 0).all():  # Hour Precision