Expand logging test cases for sample_prior_predictive and add return type overloads

pymc/sampling/forward.py

@@ -20,8 +20,10 @@
 from collections.abc import Callable, Iterable, Sequence
 from typing import (
     Any,
+    Literal,
     TypeAlias,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -360,6 +362,28 @@ def observed_dependent_deterministics(model: Model, extra_observeds=None):
     ]
 
 
+@overload
+def sample_prior_predictive(
+    draws: int = 500,
+    model: Model | None = None,
+    var_names: Iterable[str] | None = None,
+    random_seed: RandomState = None,
+    return_inferencedata: Literal[True] = True,
+    idata_kwargs: dict | None = None,
+    compile_kwargs: dict | None = None,
+    samples: int | None = None,
+) -> InferenceData: ...
+@overload
+def sample_prior_predictive(
+    draws: int = 500,
+    model: Model | None = None,
+    var_names: Iterable[str] | None = None,
+    random_seed: RandomState = None,
+    return_inferencedata: Literal[False] = False,
+    idata_kwargs: dict | None = None,
+    compile_kwargs: dict | None = None,
+    samples: int | None = None,
+) -> dict[str, np.ndarray]: ...
 def sample_prior_predictive(
     draws: int = 500,
     model: Model | None = None,
@@ -449,7 +473,9 @@ def sample_prior_predictive(
     )
 
     # All model variables have a name, but mypy does not know this
-    _log.info(f"Sampling: {sorted(volatile_basic_rvs, key=lambda var: var.name)}")  # type: ignore[arg-type, return-value]
+    _log.info(
+        f"Sampling: {sorted(volatile_basic_rvs.intersection(vars_to_sample), key=lambda var: var.name)}"  # type: ignore[arg-type, return-value]
+    )
     values = zip(*(sampler_fn() for i in range(draws)))
 
     data = {k: np.stack(v) for k, v in zip(names, values)}

tests/sampling/test_forward.py

@@ -862,11 +862,18 @@ def test_logging_sampled_basic_rvs_prior(self, caplog):
         assert caplog.record_tuples == [("pymc.sampling.forward", logging.INFO, "Sampling: [x, z]")]
         caplog.clear()
 
+        # RV only
         with m:
             pm.sample_prior_predictive(draws=1, var_names=["x"])
         assert caplog.record_tuples == [("pymc.sampling.forward", logging.INFO, "Sampling: [x]")]
         caplog.clear()
 
+        # observed only
+        with m:
+            pm.sample_prior_predictive(draws=1, var_names=["z"])
+        assert caplog.record_tuples == [("pymc.sampling.forward", logging.INFO, "Sampling: [z]")]
+        caplog.clear()
+
     def test_logging_sampled_basic_rvs_posterior(self, caplog):
         with pm.Model() as m:
             x = pm.Normal("x")