Narrow ImportError suppression, fix Bacon docs, and fix plot return types

Restrict test harness ImportError handling to known third-party modules
(pyfixest, linearmodels, differences) so broken diff_diff imports fail CI.
Fix REGISTRY.md Bacon never-treated assumption to match implementation.
Fix fig= to ax= across visualization docs (plot_* returns Axes, not Figure).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

docs/api/power.rst

@@ -157,8 +157,8 @@ Complete Example
 
    # Power curve
    pa = PowerAnalysis(n_treated=100, n_control=100, n_pre=4, n_post=4, sigma=1.0)
-   fig = plot_power_curve(pa, effect_range=(0, 1), n_points=50)
-   fig.savefig('power_curve.png')
+   ax = plot_power_curve(pa, effect_range=(0, 1), n_points=50)
+   ax.figure.savefig('power_curve.png')
 
 See Also
 --------

docs/api/visualization.rst

@@ -47,7 +47,7 @@ Example
                     time='period', first_treat='first_treat')
 
    # Plot effects by treatment cohort
-   fig = plot_group_effects(results)
+   ax = plot_group_effects(results)
 
 plot_sensitivity
 ----------------
@@ -69,7 +69,7 @@ Example
        M_grid=[0, 0.5, 1.0, 1.5, 2.0]
    )
 
-   fig = plot_sensitivity(sensitivity)
+   ax = plot_sensitivity(sensitivity)
 
 plot_honest_event_study
 -----------------------

docs/methodology/REGISTRY.md

@@ -1500,7 +1500,7 @@ For global method, LOOCV works as follows:
 *Assumption checks / warnings:*
 - Requires variation in treatment timing (staggered adoption)
 - Warns if only one treatment cohort (decomposition not meaningful)
-- Assumes no never-treated: uses not-yet-treated as controls
+- Uses never-treated units as controls when present; falls back to timing-only comparisons otherwise
 
 *Estimator equation (as implemented):*
 

docs/troubleshooting.rst

@@ -252,15 +252,15 @@ Visualization Issues
    import matplotlib.pyplot as plt
 
    # Option 1: Use plt.show()
-   fig = plot_event_study(results)
+   ax = plot_event_study(results)
    plt.show()
 
    # Option 2: Use inline magic (Jupyter)
    %matplotlib inline
 
    # Option 3: Return and display figure
-   fig = plot_event_study(results)
-   fig  # Display in Jupyter
+   ax = plot_event_study(results)
+   ax  # Display in Jupyter
 
 Performance Issues
 ------------------

tests/test_doc_snippets.py

@@ -3,7 +3,8 @@
 
 Extracts ``.. code-block:: python`` snippets from RST files and executes them
 in isolated namespaces with synthetic data and mock dataset loaders. Fails on
-all exceptions except NameError (context-dependent snippets).
+all exceptions except NameError (context-dependent snippets) and
+ImportError for known third-party packages (comparison-page snippets).
 """
 
 import re
@@ -100,6 +101,11 @@ def _extract_snippets(rst_path: Path) -> List[Tuple[int, str]]:
     r"wide_to_long\(",  # references undefined wide_data variable
 ]
 
+# Third-party packages imported by comparison-page snippets that may not
+# be installed in the test environment.  Only these are exempt from
+# ImportError failures — diff_diff and stdlib imports must succeed.
+_THIRD_PARTY_MODULES = {"pyfixest", "linearmodels", "differences"}
+
 
 def _should_skip(code: str) -> Optional[str]:
     """Return a reason string if the snippet should be skipped, else None."""
@@ -354,11 +360,18 @@ def test_doc_snippet(test_id: str, code: str, skip_reason: Optional[str]):
         # context block (e.g. ``results`` from an earlier fit).  This is
         # expected for isolated execution — not an API mismatch.
         pass
-    except ImportError:
-        # ImportError covers both ModuleNotFoundError (comparison pages
-        # importing pyfixest, linearmodels, etc.) and optional-dependency
-        # guards (e.g. matplotlib required for plotting functions).
-        pass
+    except ImportError as exc:
+        # Only suppress ImportError for known third-party packages that
+        # comparison-page snippets import. In-package (diff_diff.*) and
+        # stdlib import failures should still fail the test.
+        mod_name = getattr(exc, "name", "") or ""
+        top_level = mod_name.split(".")[0]
+        if top_level not in _THIRD_PARTY_MODULES:
+            pytest.fail(
+                f"Snippet {test_id} raised ImportError for "
+                f"'{mod_name}': {exc}\n\n"
+                f"Code:\n{textwrap.indent(code, '  ')}"
+            )
     except Exception as exc:
         pytest.fail(
             f"Snippet {test_id} raised {type(exc).__name__}: {exc}\n\n"