Add support for ArviZ <1.0, on top of existing >1.0 support (#22)

* Support both xr.DataTree and arviz.InferenceData via duck-typing

Add _get_group() and _has_group() private helpers that dispatch on
hasattr(dt, 'children') to support both arviz >=1.0 DataTree and
arviz <1.0 InferenceData with zero public API changes.

- Rename _datatree_group_to_df → _group_to_df (no longer datatree-specific)
- Route all group access through the two helpers
- Loosen arviz pin from >=1.0 to any version
- Update docstrings and error messages for dual-type support

Closes #21

* Add backend dispatch tests covering both DataTree and InferenceData paths

Parametrized tests verify parameter_draws, compare_draws, and
prediction_draws work identically with xr.DataTree and a duck-type
InferenceData mock (attribute-access groups, no .children).

42 new tests, 100% pass rate against both backends.

* Add arviz<1.0 CI matrix and real-InferenceData fixture

The legacy az.InferenceData dispatch branch was only exercised by the
_InferenceDataMock duck-type; on arviz>=1.0 az.InferenceData aliases to
DataTree so the getattr(dt, group) path was unreachable in CI.

- ci.yml: matrix.arviz=[locked, legacy]; legacy job pins arviz<1.0
  and runs tests with --no-sync so the downgrade is not reverted.
- test_backend_compat.py: add real_inferencedata variant to parameter,
  compare, and prediction fixtures. Skips on arviz>=1.0 (DetectionWarning
  path), runs same assertions against real az.InferenceData on arviz<1.0.

Verified: 92 passed/21 skipped on arviz 1.2.0; 63 passed/0 skipped on
arviz 0.23.4 (real legacy branch exercised).

* docs: add GitHub install instructions for pre-release testing

Lets users test main-branch functionality before a PyPI release:
  pip install git+https://github.com/drbenvincent/tidydraws.git
  uv add git+https://github.com/drbenvincent/tidydraws.git

Added to both README.md and index.qmd install sections.

* Address review non-blocking items: type annotations, pin, error msg, docs

- _extract.py: widen dt annotations from xr.DataTree to a PEP 695
  ArviZData = xr.DataTree | az.InferenceData union alias. az is imported
  only under TYPE_CHECKING; the alias rhs is lazy so no runtime arviz
  import is needed (matches the PR's no-import design). Applied to all
  three public functions and the _get_group/_has_group/_group_to_df helpers.
- pyproject.toml: pin arviz>=0.12 (was bare 'arviz'), matching the PR's
  stated out-of-scope (arviz<0.12 not supported).
- _extract.py: restore 'constant_data_group' parameter name in the
  prediction_draws missing-data error message (AGENTS.md Hard Rule 3).
- 04-showcase.qmd: fix stale _datatree_group_to_df -> _group_to_df ref.

Verified: 92 passed/21 skipped; pre-commit (ruff, mypy v1.15.0, format)
passes; arviz not imported at runtime; ArviZData is a lazy TypeAliasType.

Author: drbenvincent

.github/workflows/ci.yml

@@ -14,7 +14,16 @@ concurrency:
 
 jobs:
   test:
+    name: test (arviz ${{ matrix.arviz }})
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        # "locked" uses the arviz version pinned in uv.lock (>=1.0).
+        # "legacy" downgrades to arviz<1.0 to exercise the real
+        # az.InferenceData dispatch branch (skipped on >=1.0 where
+        # az.InferenceData is aliased to xarray.DataTree).
+        arviz: [locked, legacy]
     permissions:
       contents: read
     steps:
@@ -25,8 +34,15 @@ jobs:
         uses: astral-sh/setup-uv@v8.2.0
       - name: Install dependencies
         run: uv sync --all-extras
+      - name: Pin legacy arviz (<1.0)
+        if: matrix.arviz == 'legacy'
+        # Override the lock's arviz>=1.0 with a legacy release so the
+        # real-InferenceData test variants activate instead of skipping.
+        run: uv pip install 'arviz<1.0'
       - name: Run tests
-        run: uv run pytest -v
+        # --no-sync so the legacy-arviz downgrade from the previous step
+        # is not reverted by `uv run` reconciling back to uv.lock.
+        run: uv run --no-sync pytest -v
 
   prek:
     runs-on: ubuntu-latest

AGENTS.md

@@ -51,7 +51,7 @@ Key point: prefix Python/pytest commands with `uv run` so they use the project's
 - **Fail loudly** on missing data to avoid silent misalignment bugs; give the user explicit guidance.
 - **Cross-dim joins auto-join with a logged warning** (e.g. scalar `sigma` broadcast across `beta[groups]`) — common case "just works", transparently.
 - **`compare_draws()` ships in v0.1** because prior vs. posterior comparison is fundamental.
-- **ArviZ 1.0 DataTree only** — greenfield, no legacy `InferenceData` debt. Access groups via `.children[group].to_dataset()`.
+- **Supports both `xr.DataTree` and `arviz.InferenceData`** — dispatch via `_get_group()` / `_has_group()` helpers using duck-typing.
 - **Polars over pandas** for faster joins and a consistent tidy-frame type across the API surface.
 
 ---
@@ -81,7 +81,7 @@ def compare_draws(
 
 ### Core helpers (in `_extract.py`)
 
-- `_datatree_group_to_df(dt, group)` → `pl.DataFrame` with chain, draw, and all coord columns.
+- `_group_to_df(dt, group)` → `pl.DataFrame` with chain, draw, and all coord columns.
 - `_align_dims(frames)` → inner-join same-dim frames; cross-join different-dim frames with a logged warning.
 - `_coerce_to_dataframe(newdata)` → `pl.DataFrame` from `pl.DataFrame` / `pd.DataFrame`.
 
@@ -91,7 +91,7 @@ def compare_draws(
 
 - Returning `pl.LazyFrame` or leaving a `.lazy()` / `.collect()` round-trip in the extraction path — the data layer is eager by design.
 - Confusing dimension names with coordinate names in xarray — the DataArray already knows its dims, so auto-detection gets this right.
-- Forgetting groups are accessed via `.children[group].to_dataset()`.
+- Accessing groups directly instead of through `_get_group()` / `_has_group()` — use the helpers for compat with both DataTree and InferenceData.
 - Running `pytest`/`python` without `uv run` (wrong environment).
 
 ---

README.md

@@ -26,6 +26,18 @@ With pip:
 pip install tidydraws
 ```
 
+If you want the latest functionality merged into main but not yet released, install directly from GitHub:
+
+```bash
+pip install git+https://github.com/drbenvincent/tidydraws.git
+```
+
+Or with uv:
+
+```bash
+uv add git+https://github.com/drbenvincent/tidydraws.git
+```
+
 ## Why tidydraws?
 
 Plotting MCMC output in Python means manually slicing xarray dimensions, iterating groups, and aligning coordinates — imperative, verbose, error-prone. R's [`tidybayes`](https://github.com/mjskay/tidybayes) solved this with a data layer that respects parameter space vs prediction space. `tidydraws` brings that to Python on Polars.

docs/examples/04-showcase.qmd

@@ -123,7 +123,7 @@ pred_summary = td.point_interval(
 
 ## 2. Varying intercepts
 
-One-dimensional `alpha[group]` with string coordinate labels.  Exercises `_datatree_group_to_df` with non-numeric coords and tests that forest-plot sorting/filtering works on string dims.
+One-dimensional `alpha[group]` with string coordinate labels.  Exercises `_group_to_df` with non-numeric coords and tests that forest-plot sorting/filtering works on string dims.
 
 ### Use tidydraws
 

index.qmd

@@ -305,21 +305,33 @@ alt_pred = (
 
 ## Install
 
-:::: {.panel-tabset}
+::::: {.panel-tabset}
 
 #### uv
 
 ```bash
 uv add tidydraws
 ```
 
+If you want the latest functionality merged into main but not yet released, install directly from GitHub:
+
+```bash
+uv add git+https://github.com/drbenvincent/tidydraws.git
+```
+
 #### pip
 
 ```bash
 pip install tidydraws
 ```
 
-::::
+If you want the latest functionality merged into main but not yet released, install directly from GitHub:
+
+```bash
+pip install git+https://github.com/drbenvincent/tidydraws.git
+```
+
+:::::
 
 ## Why tidydraws?
 

pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
 requires-python = ">=3.12"
 dependencies = [
     "polars >=0.20",
-    "arviz >=1.0",
+    "arviz>=0.12",
     "xarray",
     "numpy",
     "pyarrow>=24.0.0",