tree-annotated-plot 0.2.0__tar.gz → 0.2.2__tar.gz

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/CHANGELOG.md

@@ -5,6 +5,38 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.2] - 2026-05-09
+
+### Fixed
+
+- `mark_line`, `mark_trail`, and `mark_area` now connect points in tree
+  tip order regardless of other encodings on the chart. Previously, a
+  user chart with an explicit categorical color-scale `domain` rendered
+  with crisscrossing line segments because Vega-Lite's default
+  connection-order heuristic ignored the strain-axis sort. The package
+  now attaches a `calculate` transform that derives a per-row tip rank
+  and points the `order` channel at that field on these marks. Any
+  user-supplied `order` is left in place.
+- When `color_tree_by` is used together with a user chart that has its
+  own `color` encoding (e.g. a titer plot colored by `cell_line`), the
+  user chart's marks no longer disappear. The concat container now
+  resolves the `color` scale as `independent` so the tree's
+  `color_value:N` scale (with its tree-specific domain) is not merged
+  with the user chart's color scale.
+
+## [0.2.1] - 2026-05-06
+
+### Fixed
+
+- Apply tree-tip ordering to chart axes that use untyped Altair
+  shorthand (`alt.Y("strain")` without a `:N` / `:O` type suffix).
+  Previously the sort override was silently skipped, so the chart
+  rendered in data order instead of tree order.
+- An internal consistency check now raises if the spec-level walk
+  and the live-object walk over the user's chart ever disagree on
+  the number of strain-axis encodings, so silent skips like the
+  one above can't recur in another shape.
+
 ## [0.2.0] - 2026-05-06
 
 ### Added

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tree-annotated-plot
-Version: 0.2.0
+Version: 0.2.2
 Summary: Annotate the axis of an Altair / Vega-Lite plot with a phylogenetic tree.
 Project-URL: Documentation, https://jbloomlab.github.io/tree-annotated-plot/
 Project-URL: Homepage, https://github.com/jbloomlab/tree-annotated-plot

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/docs/examples.md

@@ -105,8 +105,12 @@ edge.
 
 The H3N2 example above is rendered with `color_tree_by="subclade"`,
 which colors the tree's branches and tip circles by the
-`node_attrs.subclade` value at each node and adds a categorical legend
-below the plot. See "Color the tree" below for the full set of options.
+`node_attrs.subclade` value at each node. The package's default places
+the categorical legend at the bottom of the combined plot; here we pass
+`tree_color_legend_format={"orient": "left"}` to push it to the left
+side instead, so it doesn't compete with the cohort-selection legend
+already sitting below the chart. See "Color the tree" below for the
+full set of options.
 
 ### Optional: connect leaders all the way to the labels
 
@@ -156,8 +160,11 @@ Colors match what you'd see on the Nextstrain view of the same tree —
 either from the JSON's palette information when the build provides it,
 or from the same default palette Auspice uses when it doesn't.
 Categories are ordered by descending frequency in both cases. Missing
-values render in gray, and the legend is drawn at the bottom of the
-combined plot.
+values render in gray. By default the legend is drawn at the bottom of
+the combined plot; the H3N2 examples in this section pass
+`tree_color_legend_format={"orient": "left"}` to move it to the left
+instead, since the bottom edge already carries the cohort-selection
+legend.
 
 The example below colors the same H3N2 chart by genotype at HA1
 site 158, which has two mutations in the tree (`N158K`, `N158D`) and
@@ -277,6 +284,7 @@ tree-annotated-plot \
     --scale-bar \
     --branch-length-units substitutions \
     --color-tree-by subclade \
+    --tree-color-legend-format '{"orient":"left"}' \
     --output examples/data/h3n2_combined.json
 ```
 
@@ -295,6 +303,7 @@ out = tree_annotated_plot.plot(
     scale_bar=True,
     branch_length_units="substitutions",
     color_tree_by="subclade",
+    tree_color_legend_format={"orient": "left"},
 )
 ```
 

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tree-annotated-plot"
-version = "0.2.0"
+version = "0.2.2"
 description = "Annotate the axis of an Altair / Vega-Lite plot with a phylogenetic tree."
 readme = "README.md"
 requires-python = ">=3.13"

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/scripts/generate_docs_assets.py

@@ -122,7 +122,11 @@ def _render_kikawa() -> None:
             # Color H3N2 by subclade so the docs SVG matches what users see
             # on Nextstrain. The Auspice JSON's meta.colorings.subclade has
             # no `scale` defined, so colors come from the default palette.
+            # Place the subclade legend on the left so it doesn't compete
+            # with the cohort-selection legend that already sits at the
+            # bottom of the chart.
             plot_kwargs["color_tree_by"] = "subclade"
+            plot_kwargs["tree_color_legend_format"] = {"orient": "left"}
         out = tree_annotated_plot.plot(
             DATA_DIR / f"flu-seqneut-2025to2026_{subtype}.json",
             chart,
@@ -155,6 +159,7 @@ def _render_kikawa() -> None:
         strain_label_font_size=9,
         shift_tree_loc=60,
         color_tree_by="subclade",
+        tree_color_legend_format={"orient": "left"},
     )
     _save_pair(out, "h3n2_combined_label_connect")
 
@@ -180,6 +185,7 @@ def _render_kikawa() -> None:
         scale_bar=True,
         branch_length_units="substitutions",
         color_tree_by="genotype:HA1:158",
+        tree_color_legend_format={"orient": "left"},
     )
     _save_pair(out, "h3n2_combined_genotype_158")
 

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/src/tree_annotated_plot/_plot.py

@@ -7,6 +7,7 @@ import json
 import math
 import re
 import warnings
+from collections.abc import Iterator
 from html.parser import HTMLParser
 from pathlib import Path
 from typing import Any, Literal
@@ -223,11 +224,37 @@ def _build(
         legend_show=config.tree_color_legend_show,
     )
 
-    new_chart = _apply_tree_order_to_chart_object(
-        chart, config.chart_strain_field, tip_names
+    new_chart = copy.deepcopy(chart)
+    suppress_axis_chrome = config.connect_leader_to_label
+    n_hits = 0
+    for ch in _iter_strain_axis_channels(new_chart, config.chart_strain_field):
+        ch.sort = list(tip_names)
+        if suppress_axis_chrome:
+            ch.axis = alt.Axis(labels=False, ticks=False, domain=False, title=None)
+        n_hits += 1
+    _check_walker_hits("strain-axis update", n_hits, len(axis_hits), axis)
+
+    # Pin line/trail/area connection order to tip order via Vega-Lite's
+    # `order` channel. Without this, an explicit categorical `color` scale
+    # `domain` (or other encoding choices) can shift Vega-Lite's
+    # connection-order heuristic away from the strain-axis sort, causing
+    # lines to crisscross. The order channel's `sort` only accepts
+    # ascending/descending, so we attach a `calculate` transform that
+    # computes a per-row tip rank and point `order` at that derived
+    # quantitative field. User-supplied `order` always wins.
+    rank_expr = (
+        f"indexof({json.dumps(list(tip_names))}, "
+        f"datum[{json.dumps(config.chart_strain_field)}])"
     )
-    if config.connect_leader_to_label:
-        _suppress_chart_strain_axis(new_chart, config.chart_strain_field)
+    for node in _iter_connection_order_nodes(new_chart, config.chart_strain_field):
+        enc = _live_attr(node, "encoding")
+        if _live_attr(enc, "order") is not None:
+            continue
+        existing = _live_attr(node, "transform") or []
+        node.transform = list(existing) + [
+            {"calculate": rank_expr, "as": _TIP_ORDER_RANK_FIELD}
+        ]
+        enc.order = alt.Order(f"{_TIP_ORDER_RANK_FIELD}:Q")
     hoisted_config, hoisted_other = _pop_toplevel_only_attrs(new_chart)
 
     combined = _concat_for_location(
@@ -271,22 +298,32 @@ def _concat_for_location(
     user_chart: alt.TopLevelMixin,
     location: TreeLocation,
 ) -> alt.HConcatChart | alt.VConcatChart:
-    """Concat tree and chart in the order implied by the tree's location."""
+    """Concat tree and chart in the order implied by the tree's location.
+
+    The strain axis is resolved independent so the tree and chart can use
+    different scales on that axis (the tree's branch length vs. the chart's
+    measurement value), while still sharing the orthogonal strain axis. The
+    `color` scale is also resolved independent: when ``color_tree_by`` is set
+    the tree panel emits a `color_value:N` color scale with a tree-specific
+    domain, and Vega-Lite's default of sharing color across concat views
+    would merge it with any color encoding on the user's chart, hiding
+    user-chart marks whose color values aren't in the tree's domain.
+    """
     if location == "left":
         return alt.hconcat(tree_chart, user_chart, spacing=0).resolve_scale(
-            y="independent"
+            y="independent", color="independent"
         )
     if location == "right":
         return alt.hconcat(user_chart, tree_chart, spacing=0).resolve_scale(
-            y="independent"
+            y="independent", color="independent"
         )
     if location == "top":
         return alt.vconcat(tree_chart, user_chart, spacing=0).resolve_scale(
-            x="independent"
+            x="independent", color="independent"
         )
     if location == "bottom":
         return alt.vconcat(user_chart, tree_chart, spacing=0).resolve_scale(
-            x="independent"
+            x="independent", color="independent"
         )
     raise ValueError(f"unreachable: tree_location={location!r}")
 
@@ -916,11 +953,45 @@ def _enumerate_tree_tip_attrs(
     return out
 
 
-def _apply_tree_order_to_chart_object(
-    chart: alt.TopLevelMixin, chart_strain_field: str, sort_order: list[str]
-) -> alt.TopLevelMixin:
-    """Return a deepcopy of the user's chart with `sort=sort_order` applied
-    to every axis encoding that references `chart_strain_field`.
+def _check_walker_hits(operation: str, actual: int, expected: int, axis: str) -> None:
+    """Cross-check the live-object walk's hit count against the spec walk's.
+
+    `_find_strain_encoding` walks the chart's serialized dict form to count
+    how many `x`/`y` strain encodings the chart has and to validate
+    consistency (axis agreement, type=nominal/ordinal, etc.). The live
+    iteration in `_build` (driven by `_iter_strain_axis_channels`) must
+    visit exactly the same number of encodings — fewer would silently skip
+    applying the tree's tip order (or the axis suppression) to part of the
+    chart, more would mean we mutated structures the dict walker doesn't
+    know about.
+
+    The check is symmetric (`!=`) rather than one-sided (`<`) because
+    either direction signals that spec-level introspection and
+    live-object traversal have diverged, and continuing would render an
+    unverified chart.
+    """
+    if actual != expected:
+        raise RuntimeError(
+            f"internal consistency check failed for {operation!r}: "
+            f"_find_strain_encoding located {expected} strain {axis!r}-axis "
+            f"encoding(s) in the chart spec, but the live-object walk "
+            f"updated {actual}. Spec-level and live-object traversal have "
+            "diverged, which would render the chart with a wrong tip order "
+            "or leave axis chrome behind. Please file a bug at "
+            "https://github.com/jbloomlab/tree-annotated-plot/issues with a "
+            "minimal reproducer."
+        )
+
+
+def _iter_strain_axis_channels(node: Any, chart_strain_field: str) -> Iterator[Any]:
+    """Yield every live x/y channel object whose field matches
+    `chart_strain_field`.
+
+    Pure read — no mutation, no count. The caller iterates and applies
+    whatever mutation it needs (currently `sort` and, when
+    `connect_leader_to_label=True`, an axis-suppression `alt.Axis(...)`),
+    counting hits as it goes so the cross-check in `_check_walker_hits`
+    can compare against the spec walker.
 
     Walks the live altair object tree (Chart / LayerChart / FacetChart /
     HConcatChart / VConcatChart / ConcatChart) rather than its dict form,
@@ -929,62 +1000,25 @@ def _apply_tree_order_to_chart_object(
     dict approach has to fight. Modifying the object in place is robust as
     long as altair's container attribute names hold (.hconcat / .vconcat /
     .concat / .layer / .spec / .encoding), which is stable in altair 5+.
+    `FacetChart.spec` is recursed into unconditionally so we descend to the
+    inner LayerChart / Chart that actually carries the encoding (gating on
+    `spec.encoding is not None` would skip LayerCharts, whose encodings
+    live on their layers rather than at the top level).
     """
-    new_chart = copy.deepcopy(chart)
-    _walk_and_apply_sort(new_chart, chart_strain_field, sort_order)
-    return new_chart
-
-
-def _walk_and_apply_sort(
-    node: Any, chart_strain_field: str, sort_order: list[str]
-) -> None:
-    """Recursively set sort on every encoding whose field == chart_strain_field
-    on a live altair chart object."""
     enc = _live_attr(node, "encoding")
     if enc is not None:
         for channel in ("x", "y"):
             ch = _live_attr(enc, channel)
             if ch is not None and _channel_field(ch) == chart_strain_field:
-                ch.sort = list(sort_order)
+                yield ch
     for attr in ("hconcat", "vconcat", "concat", "layer"):
         sub = _live_attr(node, attr)
         if isinstance(sub, list):
             for s in sub:
-                _walk_and_apply_sort(s, chart_strain_field, sort_order)
-    # FacetChart.spec is the chart being faceted; recurse unconditionally so
-    # we descend into the inner LayerChart / Chart that actually carries the
-    # encoding. Gating on `spec.encoding is not None` was wrong because a
-    # LayerChart has no top-level encoding — its encodings live on its layers.
+                yield from _iter_strain_axis_channels(s, chart_strain_field)
     spec = _live_attr(node, "spec")
     if spec is not None:
-        _walk_and_apply_sort(spec, chart_strain_field, sort_order)
-
-
-def _suppress_chart_strain_axis(
-    chart: alt.TopLevelMixin, chart_strain_field: str
-) -> None:
-    """Hide labels, ticks, axis line, and title on every chart strain-axis encoding.
-
-    Walks the live altair object the same way `_walk_and_apply_sort` does
-    and replaces the matching encoding's ``axis`` with one that suppresses
-    every visible bit of axis chrome. This **overrides** any user-supplied
-    ``axis=alt.Axis(...)`` on those encodings; that's documented on the
-    ``connect_leader_to_label`` description in `_config.py`.
-    """
-    enc = _live_attr(chart, "encoding")
-    if enc is not None:
-        for channel in ("x", "y"):
-            ch = _live_attr(enc, channel)
-            if ch is not None and _channel_field(ch) == chart_strain_field:
-                ch.axis = alt.Axis(labels=False, ticks=False, domain=False, title=None)
-    for attr in ("hconcat", "vconcat", "concat", "layer"):
-        sub = _live_attr(chart, attr)
-        if isinstance(sub, list):
-            for s in sub:
-                _suppress_chart_strain_axis(s, chart_strain_field)
-    spec = _live_attr(chart, "spec")
-    if spec is not None:
-        _suppress_chart_strain_axis(spec, chart_strain_field)
+        yield from _iter_strain_axis_channels(spec, chart_strain_field)
 
 
 def _live_attr(obj: Any, name: str) -> Any:
@@ -997,6 +1031,65 @@ def _live_attr(obj: Any, name: str) -> Any:
     return v
 
 
+# Vega-Lite marks whose drawing order along the discrete axis is determined
+# by the order channel (rule 1 of the connection-order fallback chain). For
+# these marks pinning the order channel to a per-row tip-order rank ensures
+# the line / trail / area connects strains in tip order regardless of how
+# Vega-Lite would otherwise resolve the fallback (e.g. when an explicit
+# color-scale `domain` shifts it away from the y-axis sort).
+_CONNECTION_ORDER_MARKS = frozenset({"line", "trail", "area"})
+
+# Derived field name appended to each connection-order mark's transform
+# pipeline. Chosen to be unlikely to collide with user data; private prefix
+# makes the intent clear if anyone inspects the rendered spec.
+_TIP_ORDER_RANK_FIELD = "_tap_strain_order_idx"
+
+
+def _mark_type(node: Any) -> str | None:
+    """Return the mark type string for a chart node, or None.
+
+    Handles altair's two mark forms: the plain string (e.g. ``"line"`` from
+    ``chart.mark_line()``) and the ``MarkDef`` object (e.g. from
+    ``chart.mark_line(point=True)``). Container nodes (LayerChart,
+    HConcatChart, etc.) have no mark and return None.
+    """
+    mark = _live_attr(node, "mark")
+    if mark is None:
+        return None
+    if isinstance(mark, str):
+        return mark
+    type_ = _live_attr(mark, "type")
+    return type_ if isinstance(type_, str) else None
+
+
+def _iter_connection_order_nodes(node: Any, chart_strain_field: str) -> Iterator[Any]:
+    """Yield the live chart node for every node whose mark is in
+    ``_CONNECTION_ORDER_MARKS`` and whose ``x``/``y`` strain encoding
+    matches ``chart_strain_field``.
+
+    Mirrors `_iter_strain_axis_channels`'s traversal (hconcat / vconcat /
+    concat / layer / spec descent). Yields the node itself (not just the
+    encoding) so the caller can attach both a ``calculate`` transform and
+    the ``order`` channel.
+    """
+    if _mark_type(node) in _CONNECTION_ORDER_MARKS:
+        enc = _live_attr(node, "encoding")
+        if enc is not None:
+            for channel in ("x", "y"):
+                ch = _live_attr(enc, channel)
+                if ch is not None and _channel_field(ch) == chart_strain_field:
+                    yield node
+                    break
+    for attr in ("hconcat", "vconcat", "concat", "layer"):
+        sub = _live_attr(node, attr)
+        if isinstance(sub, list):
+            for s in sub:
+                yield from _iter_connection_order_nodes(s, chart_strain_field)
+    spec = _live_attr(node, "spec")
+    if spec is not None:
+        yield from _iter_connection_order_nodes(spec, chart_strain_field)
+
+
 def _pop_toplevel_only_attrs(
     chart: alt.TopLevelMixin,
 ) -> tuple[Any, dict]:
@@ -1044,21 +1137,71 @@ def _apply_combined_config(combined: alt.HConcatChart, hoisted_config: Any) -> N
 def _channel_field(ch: Any) -> str | None:
     """Read the underlying `field` string from an altair channel object.
 
-    `ch.field` returns altair's `_PropertySetter` (used for fluent chaining),
-    not the stored field name. The stored value is reachable via `to_dict()`.
+    Returns the field name when the channel references a data field (either
+    via the `field=` keyword or via positional shorthand like
+    `alt.Y("strain")` / `alt.Y("strain:N")` / `alt.Y("mean(titer):Q")`).
+    Returns `None` when the channel has no field at all (a `value=` /
+    `datum=` constant encoding). Raises `ValueError` when the channel's
+    `_kwds` shape is unrecognized — silent fallthrough hid a real bug
+    where untyped shorthand axes were never reordered to match the tree.
+
+    Reads `ch._kwds` directly rather than going through `ch.to_dict()`:
+    altair's `to_dict()` on a bare channel raises when the shorthand has
+    no explicit type (e.g. `alt.Y("strain")`), because the `nominal` /
+    `quantitative` inference needs the chart's data context. That
+    exception is what the previous catch-all hid.
     """
-    to_dict = getattr(ch, "to_dict", None)
-    if not callable(to_dict):
-        return None
-    try:
-        d = to_dict()
-    except Exception:
+    kwds = getattr(ch, "_kwds", None)
+    if not isinstance(kwds, dict):
+        raise ValueError(
+            f"channel object {type(ch).__name__} has no `_kwds` mapping; "
+            "this isn't a recognized altair channel encoding."
+        )
+    field = kwds.get("field")
+    if field is not None and field is not alt.Undefined:
+        # `from_dict`-roundtripped channels store the field as a
+        # `FieldName(SchemaBase)` wrapper rather than a plain str; its
+        # `to_dict()` returns the raw string, while `str(...)` returns the
+        # repr `FieldName('x')`. Cover both.
+        if isinstance(field, str):
+            if field:
+                return field
+        elif hasattr(field, "to_dict"):
+            unwrapped = field.to_dict()
+            if isinstance(unwrapped, str) and unwrapped:
+                return unwrapped
+            raise ValueError(
+                f"channel field wrapper {type(field).__name__} unwrapped "
+                f"to {unwrapped!r}; expected a non-empty string."
+            )
+        else:
+            raise ValueError(
+                f"channel field has unexpected value {field!r} "
+                f"(type {type(field).__name__}); expected a string."
+            )
+    shorthand = kwds.get("shorthand")
+    if shorthand is None or shorthand is alt.Undefined:
         return None
-    if isinstance(d, dict):
-        f = d.get("field")
-        if isinstance(f, str):
-            return f
-    return None
+    if not isinstance(shorthand, str) or not shorthand:
+        raise ValueError(
+            f"channel shorthand has unexpected value {shorthand!r}; "
+            "expected a string like 'strain', 'strain:N', or "
+            "'mean(strain):Q'."
+        )
+    # Shorthand grammar: '[<aggregate>(]<field>[)][:<type>]'.
+    bare = shorthand.split(":", 1)[0]
+    if "(" in bare:
+        if not bare.endswith(")"):
+            raise ValueError(
+                f"channel shorthand {shorthand!r} has an unbalanced "
+                "aggregate wrapper; expected 'aggregate(field)[:type]'."
+            )
+        bare = bare[bare.index("(") + 1 : -1]
+    if not bare:
+        raise ValueError(
+            f"channel shorthand {shorthand!r} parsed to an empty field name."
+        )
+    return bare
 
 
 def _chart_strain_dim(

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/tests/test_color_tree.py

@@ -663,6 +663,40 @@ def test_color_tree_by_legend_hides_internal_only_unknown():
         assert "unknown" in enc["scale"]["domain"]
 
 
+def test_color_tree_by_color_scale_independent_from_user_chart():
+    """When the tree is colored, the concat container must resolve the color
+    scale as `independent` so the tree's `color_value:N` scale doesn't merge
+    with a user-chart color encoding (which would hide user marks whose
+    color values aren't in the tree's domain)."""
+    df = pd.DataFrame(
+        {
+            "strain": ["A", "B", "C", "D"] * 2,
+            "titer": [1.0, 2.0, 4.0, 8.0, 1.5, 2.5, 4.5, 8.5],
+            "cell_line": ["X", "X", "X", "X", "Y", "Y", "Y", "Y"],
+        }
+    )
+    user_chart = (
+        alt.Chart(df)
+        .mark_line(point=True)
+        .encode(
+            x="titer:Q",
+            y=alt.Y("strain:N"),
+            color=alt.Color("cell_line:N"),
+        )
+        .properties(width=200, height=200)
+    )
+    out = tree_annotated_plot.plot(
+        _attr_auspice(),
+        user_chart,
+        **_kw(),
+        color_tree_by="subclade",
+    )
+    spec = out.to_dict()
+    resolve = spec.get("resolve") or {}
+    scale = resolve.get("scale") or {}
+    assert scale.get("color") == "independent"
+
+
 # -----------------------------------------------------------------------------
 # CLI
 # -----------------------------------------------------------------------------

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/tests/test_introspection.py

@@ -15,7 +15,7 @@ import altair as alt
 import pandas as pd
 import pytest
 
-from tree_annotated_plot._plot import _find_strain_encoding
+from tree_annotated_plot._plot import _channel_field, _find_strain_encoding
 
 DATA_DIR = Path(__file__).resolve().parent.parent / "examples" / "data"
 
@@ -139,3 +139,70 @@ def test_find_strain_encoding_quantitative_type_raises() -> None:
     spec = _flat_chart_spec(typ="quantitative")
     with pytest.raises(ValueError, match="type='quantitative'"):
         _find_strain_encoding(spec, "strain")
+
+
+# ---------- _channel_field: covers the "untyped shorthand" silent-skip bug.
+
+
+def test_channel_field_typed_shorthand() -> None:
+    """`alt.Y('strain:N')` — the form every existing test uses."""
+    ch = alt.Y("strain:N")
+    assert _channel_field(ch) == "strain"
+
+
+def test_channel_field_untyped_shorthand() -> None:
+    """`alt.Y('strain')` without a type — the form that triggered the
+    original silent-skip bug. `to_dict()` on this channel raises because
+    type inference needs the chart's data, but we can still recover the
+    field from `_kwds['shorthand']`."""
+    ch = alt.Y("strain")
+    assert _channel_field(ch) == "strain"
+
+
+def test_channel_field_explicit_field_kwarg() -> None:
+    """`alt.Y(field='strain', type='nominal')` — the explicit form."""
+    ch = alt.Y(field="strain", type="nominal")
+    assert _channel_field(ch) == "strain"
+
+
+def test_channel_field_after_from_dict_roundtrip() -> None:
+    """After `alt.Chart.from_dict(...)` (used by the CLI / JSON / HTML
+    chart loaders), the field is stored as a `FieldName(SchemaBase)`
+    wrapper, not a plain `str`. The helper must unwrap it."""
+    df = pd.DataFrame({"strain": ["A", "B"], "titer": [1.0, 2.0]})
+    chart = alt.Chart(df).mark_circle().encode(x="titer:Q", y=alt.Y("strain:N"))
+    roundtripped = alt.Chart.from_dict(chart.to_dict())
+    assert _channel_field(roundtripped.encoding.y) == "strain"
+
+
+def test_channel_field_aggregate_shorthand() -> None:
+    """`alt.X('mean(titer):Q')` aggregate-wrapped shorthand."""
+    ch = alt.X("mean(titer):Q")
+    assert _channel_field(ch) == "titer"
+
+
+def test_channel_field_value_only_returns_none() -> None:
+    """A `value=` constant encoding has no field — legitimate None,
+    not an error."""
+    ch = alt.Y(value=5)
+    assert _channel_field(ch) is None
+
+
+def test_channel_field_unbalanced_aggregate_raises() -> None:
+    """A shorthand string we can't parse must raise rather than silently
+    return None — silent fallthrough is what hid the original bug."""
+    ch = alt.Y("strain")
+    ch._kwds["shorthand"] = "mean(titer:Q"
+    with pytest.raises(ValueError, match="unbalanced aggregate"):
+        _channel_field(ch)
+
+
+def test_channel_field_non_kwds_object_raises() -> None:
+    """A non-altair object accidentally passed in must raise, not return
+    None — same fail-fast principle."""
+
+    class NotAChannel:
+        pass
+
+    with pytest.raises(ValueError, match="no `_kwds` mapping"):
+        _channel_field(NotAChannel())

tree_annotated_plot-0.2.2/tests/test_line_order.py

@@ -0,0 +1,258 @@
+"""Connection order on `mark_line` / `mark_trail` / `mark_area` follows tip order.
+
+Regression: when the user chart had `mark_line` with strain on a discrete
+axis, lines connected points in tip order most of the time — but adding an
+explicit categorical color-scale `domain` shifted Vega-Lite's default
+connection-order heuristic away from the strain-axis `sort`, causing lines
+to crisscross. The package now sets the `order` channel to a per-row
+tip-rank field on these marks (unless the user has set their own `order`),
+pinning rule 1 of the Vega-Lite fallback chain.
+
+The tip-rank is computed via a `calculate` transform that runs `indexof`
+against the tip-names list — Vega-Lite's `order` channel `sort` only
+accepts ascending/descending, so a derived quantitative field is the
+direct way to express custom ordering.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import altair as alt
+import pandas as pd
+import pytest
+
+import tree_annotated_plot
+from tree_annotated_plot._plot import _TIP_ORDER_RANK_FIELD
+
+
+def _synthetic_auspice() -> dict:
+    """4-tip Auspice tree; tip layout order is [A, B, C, D]."""
+    return {
+        "version": "v2",
+        "meta": {},
+        "tree": {
+            "name": "ROOT",
+            "node_attrs": {"div": 0.0},
+            "children": [
+                {
+                    "name": "int1",
+                    "node_attrs": {"div": 0.01},
+                    "children": [
+                        {"name": "A", "node_attrs": {"div": 0.04}},
+                        {"name": "B", "node_attrs": {"div": 0.05}},
+                    ],
+                },
+                {
+                    "name": "int2",
+                    "node_attrs": {"div": 0.02},
+                    "children": [
+                        {"name": "C", "node_attrs": {"div": 0.03}},
+                        {"name": "D", "node_attrs": {"div": 0.06}},
+                    ],
+                },
+            ],
+        },
+    }
+
+
+TIP_ORDER = ["A", "B", "C", "D"]
+
+
+def _df() -> pd.DataFrame:
+    return pd.DataFrame(
+        [
+            {"strain": s, "serum": serum, "titer": v}
+            for s, serum, v in [
+                ("A", "s1", 100),
+                ("B", "s1", 200),
+                ("C", "s1", 400),
+                ("D", "s1", 800),
+                ("A", "s2", 150),
+                ("B", "s2", 250),
+                ("C", "s2", 350),
+                ("D", "s2", 700),
+            ]
+        ]
+    )
+
+
+def _plot_user_panel(chart: alt.TopLevelMixin) -> dict:
+    """Run plot() with default settings and return the user-panel spec dict."""
+    out = tree_annotated_plot.plot(
+        _synthetic_auspice(),
+        chart,
+        chart_strain_field="strain",
+        tree_strain_field="name",
+        branch_length="div",
+    )
+    # Default tree_location for y-axis strain is "left" — user panel is hconcat[1].
+    return out.to_dict()["hconcat"][1]
+
+
+def _find_order_encodings(spec: Any) -> list[dict]:
+    """Walk a spec dict and return every `encoding.order` dict found."""
+    found: list[dict] = []
+
+    def walk(node: Any) -> None:
+        if isinstance(node, dict):
+            enc = node.get("encoding")
+            if isinstance(enc, dict) and "order" in enc:
+                found.append(enc["order"])
+            for k, v in node.items():
+                if k != "encoding":
+                    walk(v)
+        elif isinstance(node, list):
+            for item in node:
+                walk(item)
+
+    walk(spec)
+    return found
+
+
+def _find_tip_rank_transforms(spec: Any) -> list[dict]:
+    """Walk a spec dict and return every calculate transform that computes
+    the package's tip-rank derived field."""
+    found: list[dict] = []
+
+    def walk(node: Any) -> None:
+        if isinstance(node, dict):
+            transforms = node.get("transform")
+            if isinstance(transforms, list):
+                for t in transforms:
+                    if isinstance(t, dict) and t.get("as") == _TIP_ORDER_RANK_FIELD:
+                        found.append(t)
+            for v in node.values():
+                walk(v)
+        elif isinstance(node, list):
+            for item in node:
+                walk(item)
+
+    walk(spec)
+    return found
+
+
+def _assert_tip_rank_wired(spec: dict) -> None:
+    """Spec contains both an order-channel pointing at the tip-rank field
+    and a calculate transform that derives that field from the tip names."""
+    orders = _find_order_encodings(spec)
+    assert len(orders) == 1, f"expected one order encoding, got {orders}"
+    assert orders[0]["field"] == _TIP_ORDER_RANK_FIELD
+    assert orders[0]["type"] == "quantitative"
+    transforms = _find_tip_rank_transforms(spec)
+    assert len(transforms) == 1, f"expected one tip-rank transform, got {transforms}"
+    expr = transforms[0]["calculate"]
+    assert "indexof" in expr
+    # The expression embeds the tip names list verbatim.
+    for tip in TIP_ORDER:
+        assert f'"{tip}"' in expr, f"tip {tip!r} not in calculate expression {expr!r}"
+
+
+def test_mark_line_with_color_domain_gets_order_in_tip_order():
+    """The original bug: explicit color-scale `domain` no longer crisscrosses lines."""
+    chart = (
+        alt.Chart(_df())
+        .mark_line(point=True)
+        .encode(
+            x=alt.X("titer:Q", scale=alt.Scale(type="log")),
+            y=alt.Y("strain:N"),
+            color=alt.Color(
+                "serum:N",
+                scale=alt.Scale(domain=["s1", "s2"], range=["red", "blue"]),
+            ),
+        )
+        .properties(width=300, height=200)
+    )
+    _assert_tip_rank_wired(_plot_user_panel(chart))
+
+
+def test_user_supplied_order_is_preserved():
+    """User-supplied `order` always wins; the package never overwrites it."""
+    chart = (
+        alt.Chart(_df())
+        .mark_line()
+        .encode(
+            x=alt.X("titer:Q"),
+            y=alt.Y("strain:N"),
+            color="serum:N",
+            order=alt.Order("titer:Q", sort="ascending"),
+        )
+        .properties(width=300, height=200)
+    )
+    spec = _plot_user_panel(chart)
+    orders = _find_order_encodings(spec)
+    assert len(orders) == 1
+    assert orders[0]["field"] == "titer"
+    # The package did not attach its tip-rank transform either.
+    assert _find_tip_rank_transforms(spec) == []
+
+
+@pytest.mark.parametrize("mark_method", ["mark_line", "mark_area", "mark_trail"])
+def test_connection_order_marks_get_order_in_tip_order(mark_method: str):
+    """All three connection-order marks (line / area / trail) get `order` injected."""
+    base = alt.Chart(_df()).encode(
+        x=alt.X("titer:Q"),
+        y=alt.Y("strain:N"),
+        color="serum:N",
+    )
+    chart = getattr(base, mark_method)().properties(width=300, height=200)
+    _assert_tip_rank_wired(_plot_user_panel(chart))
+
+
+def test_layered_chart_only_line_layer_gets_order():
+    """Mixed-mark layers: line layer gets `order`, circle layer is untouched."""
+    base = alt.Chart(_df()).encode(
+        x=alt.X("titer:Q"),
+        y=alt.Y("strain:N"),
+        color="serum:N",
+    )
+    layered = (base.mark_line() + base.mark_circle()).properties(width=300, height=200)
+    _assert_tip_rank_wired(_plot_user_panel(layered))
+
+
+def test_no_connection_order_marks_means_no_order_injected():
+    """A chart with only point/circle/etc. marks is a no-op for the walker."""
+    chart = (
+        alt.Chart(_df())
+        .mark_circle()
+        .encode(
+            x=alt.X("titer:Q"),
+            y=alt.Y("strain:N"),
+            color="serum:N",
+        )
+        .properties(width=300, height=200)
+    )
+    spec = _plot_user_panel(chart)
+    assert _find_order_encodings(spec) == []
+    assert _find_tip_rank_transforms(spec) == []
+
+
+def test_faceted_chart_gets_order_injection():
+    """The walker descends into the inner spec of a faceted chart."""
+    df_facet = pd.DataFrame(
+        [
+            {"strain": s, "serum": serum, "panel": p, "titer": v}
+            for s, serum, p, v in [
+                ("A", "s1", "P1", 100),
+                ("B", "s1", "P1", 200),
+                ("C", "s1", "P1", 400),
+                ("D", "s1", "P1", 800),
+                ("A", "s1", "P2", 150),
+                ("B", "s1", "P2", 250),
+                ("C", "s1", "P2", 350),
+                ("D", "s1", "P2", 700),
+            ]
+        ]
+    )
+    base = (
+        alt.Chart(df_facet)
+        .mark_line()
+        .encode(
+            x=alt.X("titer:Q"),
+            y=alt.Y("strain:N"),
+            color="serum:N",
+        )
+        .properties(width=300, height=200)
+    )
+    faceted = base.facet(facet="panel:N", columns=2)
+    _assert_tip_rank_wired(_plot_user_panel(faceted))

{tree_annotated_plot-0.2.0 → tree_annotated_plot-0.2.2}/tests/test_plot.py

@@ -110,6 +110,73 @@ def test_plot_overrides_y_sort_to_tree_tip_order():
     assert user_panel_spec["encoding"]["y"]["sort"] == ["A", "B", "C", "D"]
 
 
+def test_plot_overrides_sort_for_untyped_shorthand_y():
+    """Regression: `alt.Y('strain')` (no `:N` suffix) used to silently
+    skip the sort override because `_channel_field` called `to_dict()`
+    on the bare channel, which raises without a chart-data context."""
+    chart = (
+        alt.Chart(_synthetic_df())
+        .mark_line(point=True)
+        .encode(
+            x=alt.X("titer:Q", scale=alt.Scale(type="log")),
+            y=alt.Y("strain"),
+            color="serum:N",
+        )
+        .properties(width=300, height=200)
+    )
+    out = tree_annotated_plot.plot(
+        _synthetic_auspice(),
+        chart,
+        chart_strain_field="strain",
+        tree_strain_field="name",
+        branch_length="div",
+    )
+    user_panel_spec = out.hconcat[1].to_dict()
+    assert user_panel_spec["encoding"]["y"]["sort"] == ["A", "B", "C", "D"]
+
+
+def test_plot_overrides_sort_for_untyped_shorthand_x():
+    """Same regression on the horizontal layout."""
+    chart = (
+        alt.Chart(_synthetic_df())
+        .mark_line(point=True)
+        .encode(
+            x=alt.X("strain"),
+            y=alt.Y("titer:Q", scale=alt.Scale(type="log")),
+            color="serum:N",
+        )
+        .properties(width=200, height=300)
+    )
+    out = tree_annotated_plot.plot(
+        _synthetic_auspice(),
+        chart,
+        chart_strain_field="strain",
+        tree_strain_field="name",
+        branch_length="div",
+    )
+    # tree_location defaults to "bottom" for x-axis strain → vconcat(user, tree)
+    user_panel_spec = out.vconcat[0].to_dict()
+    assert user_panel_spec["encoding"]["x"]["sort"] == ["A", "B", "C", "D"]
+
+
+def test_plot_consistency_tripwire_fires_when_live_walker_misses(monkeypatch):
+    """If `_channel_field` ever reverts to silently failing, the spec walker
+    finds a strain encoding the live walker doesn't, and the count-based
+    consistency check raises rather than letting a wrong-ordered chart
+    render."""
+    from tree_annotated_plot import _plot
+
+    monkeypatch.setattr(_plot, "_channel_field", lambda ch: None)
+    with pytest.raises(RuntimeError, match="internal consistency check failed"):
+        tree_annotated_plot.plot(
+            _synthetic_auspice(),
+            _synthetic_chart(),
+            chart_strain_field="strain",
+            tree_strain_field="name",
+            branch_length="div",
+        )
+
+
 def test_plot_strain_mismatch_raises():
     bad = _synthetic_chart()
     bad.data = bad.data.replace({"strain": {"D": "X"}})