tree-annotated-plot 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

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.dist-info → tree_annotated_plot-0.2.2.dist-info}/METADATA

@@ -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.dist-info → tree_annotated_plot-0.2.2.dist-info}/RECORD

@@ -1,11 +1,11 @@
 tree_annotated_plot/__init__.py,sha256=R39gjVnl3K2VjLba6zwtwPGADFmtNOE2UYNk1y9_z6k,351
 tree_annotated_plot/_color.py,sha256=nOZp-KIPy6s-6gBkzW36oHrwDk_SmvVNsMiIJ-dY1Io,26519
 tree_annotated_plot/_config.py,sha256=XwokRKY-ytNXelxwuapRKedTBaEedDi_erUp3su7zUs,12498
-tree_annotated_plot/_plot.py,sha256=T49hL5giVsTJgvA6SuA7fs7M1da48NSe7unQf7zRVT8,65397
+tree_annotated_plot/_plot.py,sha256=GCdcPyfTlem9YTlPkgdkJGPb5g4w2Vt0fevJcHEo9HI,72720
 tree_annotated_plot/_tree.py,sha256=Bj-ABt8aooSwV7wAVTbwFGMmCnTPGOKweUWAdfb7aEI,13563
 tree_annotated_plot/cli.py,sha256=yXI5fcLoIoLf6EjgPLmIQIAzDfkVoS5mSZ26_HSHVmk,9786
-tree_annotated_plot-0.2.0.dist-info/METADATA,sha256=4Ib5mqP-YUKH6JhdnMwnzoY__0MwDnArMIB_mb75V1g,11353
-tree_annotated_plot-0.2.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-tree_annotated_plot-0.2.0.dist-info/entry_points.txt,sha256=K7edTocRbwmtM2OeZmCy3DLjmY3MR-q2EJmQQq4rj0E,69
-tree_annotated_plot-0.2.0.dist-info/licenses/LICENSE,sha256=UIp30ue1HZit7tMmSK0SS3swJBqOcLLnLJn5LY5YdNE,1036
-tree_annotated_plot-0.2.0.dist-info/RECORD,,
+tree_annotated_plot-0.2.2.dist-info/METADATA,sha256=8yAZMK69Tos2TmV-_Gb359Xm6OlmC8Elzzgzdp3yM04,11353
+tree_annotated_plot-0.2.2.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+tree_annotated_plot-0.2.2.dist-info/entry_points.txt,sha256=K7edTocRbwmtM2OeZmCy3DLjmY3MR-q2EJmQQq4rj0E,69
+tree_annotated_plot-0.2.2.dist-info/licenses/LICENSE,sha256=UIp30ue1HZit7tMmSK0SS3swJBqOcLLnLJn5LY5YdNE,1036
+tree_annotated_plot-0.2.2.dist-info/RECORD,,