codefig 0.17.0__py3-none-any.whl

codefig/tree.py

@@ -0,0 +1,277 @@
+"""Tidy trees (v0.4.0): ``Tree(root, child, Tree(sub, leaf, leaf), …)``.
+
+Nodes are any primitives (Box, Circle, Icon, Image, Group …). The layout is a
+Reingold–Tilford style tidy tree: subtrees are packed as close as their *contours*
+allow (not their bounding boxes), parents are centred over their children, and each
+depth level gets one row (``direction="down"``) or column (``"right"``, also "up" /
+"left"). Edges are real :class:`~codefig.edges.Fork` buses (``edges="ortho"``, the
+TikZ / draw.io tree look) or straight :class:`~codefig.edges.Edge` lines between the
+true outlines (``edges="straight"``), or none. A Tree is a :class:`Diagram`, so it nests
+anywhere and other edges can reach its nodes.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+from .edges import Diagram, Edge, Fork
+from .node import Node
+
+_DIRS = ("down", "up", "right", "left")
+
+
+class _Spec:
+    __slots__ = ("node", "kids", "edge_kw")
+
+    def __init__(self, node: Node, kids: List["_Spec"], edge_kw: Optional[Dict[str, Any]] = None):
+        self.node = node
+        self.kids = kids
+        self.edge_kw = edge_kw or {}
+
+
+def _as_spec(item: Any) -> _Spec:
+    if isinstance(item, Tree):
+        return item._spec
+    if isinstance(item, (tuple, list)):
+        if not item:
+            raise ValueError("empty subtree spec")
+        return _Spec(item[0], [_as_spec(k) for k in item[1:]])
+    if isinstance(item, Node):
+        return _Spec(item, [])
+    raise TypeError(f"Tree children must be nodes, Trees or (node, *children) tuples, got {type(item).__name__}")
+
+
+class _TreeLayout(Node):
+    """Layout-only content of a Tree: places every node; draws them (no frame)."""
+
+    def __init__(self, tree: "Tree", theme: Any = None) -> None:
+        super().__init__(theme=theme)
+        self.tree = tree
+        self._nodes: List[Node] = []
+        for s, _p in tree._walk():
+            self._nodes.append(s.node)
+            s.node.parent = self
+        self._pos: Dict[int, Tuple[float, float]] = {}  # node id → (breadth centre, depth start)
+        self._sizes: Dict[int, Tuple[float, float]] = {}
+
+    def children(self) -> Sequence[Node]:
+        return self._nodes
+
+    def invalidate(self) -> None:
+        super().invalidate()
+        for n in self._nodes:
+            n.invalidate()
+
+    def _bd(self, n: Node) -> Tuple[float, float]:
+        w, h = n.measure()
+        return (w, h) if self.tree.direction in ("down", "up") else (h, w)
+
+    def _measure(self) -> Tuple[float, float]:
+        t = self.tree
+        gap, lgap = t.gap, t.level_gap
+        # depth thickness per level
+        levels: Dict[int, float] = {}
+
+        def walk(s: _Spec, d: int) -> None:
+            b, dd = self._bd(s.node)
+            levels[d] = max(levels.get(d, 0.0), dd)
+            for k in s.kids:
+                walk(k, d + 1)
+        walk(t._spec, 0)
+
+        offs: Dict[int, float] = {}
+
+        def lay(s: _Spec, d: int) -> Tuple[Dict[int, float], Dict[int, Tuple[float, float]]]:
+            b, _ = self._bd(s.node)
+            pos: Dict[int, float] = {id(s.node): 0.0}
+            cont: Dict[int, Tuple[float, float]] = {d: (-b / 2, b / 2)}
+            if not s.kids:
+                return pos, cont
+            merged: Dict[int, Tuple[float, float]] = {}
+            placed: List[Tuple[float, Dict[int, float]]] = []
+            for k in s.kids:
+                kp, kc = lay(k, d + 1)
+                if not merged:
+                    o = 0.0
+                else:
+                    o = max(merged[dd][1] - kc[dd][0] + gap for dd in kc if dd in merged) \
+                        if any(dd in merged for dd in kc) else 0.0
+                for dd, (l, r) in kc.items():
+                    if dd in merged:
+                        merged[dd] = (min(merged[dd][0], l + o), max(merged[dd][1], r + o))
+                    else:
+                        merged[dd] = (l + o, r + o)
+                placed.append((o, kp))
+            first, last = placed[0][0], placed[-1][0]
+            if t.align == "first":
+                c = first
+            else:
+                c = (first + last) / 2
+            for o, kp in placed:
+                for nid, v in kp.items():
+                    pos[nid] = v + o - c
+            for dd, (l, r) in merged.items():
+                cont[dd] = (l - c, r - c) if dd != d else cont[dd]
+            return pos, cont
+
+        pos, cont = lay(t._spec, 0)
+        lo = min(l for l, _ in cont.values())
+        hi = max(r for _, r in cont.values())
+        depth_start: Dict[int, float] = {}
+        acc = 0.0
+        for d in sorted(levels):
+            depth_start[d] = acc
+            acc += levels[d] + lgap
+        total_depth = acc - lgap
+
+        self._nodes = []
+        self._pos = {}
+
+        def collect(s: _Spec, d: int) -> None:
+            self._nodes.append(s.node)
+            s.node.parent = self
+            b, dd = self._bd(s.node)
+            self._pos[id(s.node)] = (pos[id(s.node)] - lo, depth_start[d] + (levels[d] - dd) / 2)
+            for k in s.kids:
+                collect(k, d + 1)
+        collect(t._spec, 0)
+        self._extent = (hi - lo, total_depth)
+        bw, dw = hi - lo, total_depth
+        return (bw, dw) if t.direction in ("down", "up") else (dw, bw)
+
+    def _layout_children(self) -> None:
+        g = self.geom
+        t = self.tree
+        bw, dw = self._extent
+        for n in self._nodes:
+            w, h = n.measure()
+            bc, ds = self._pos[id(n)]
+            b, dd = self._bd(n)
+            if t.direction == "down":
+                n.layout(g.x + bc - w / 2, g.y + ds, w, h)
+            elif t.direction == "up":
+                n.layout(g.x + bc - w / 2, g.y + dw - ds - h, w, h)
+            elif t.direction == "right":
+                n.layout(g.x + ds, g.y + bc - h / 2, w, h)
+            else:
+                n.layout(g.x + dw - ds - w, g.y + bc - h / 2, w, h)
+
+    def draw(self, buf: List[str], ox: float = 0.0, oy: float = 0.0) -> None:
+        for n in self._nodes:
+            n.draw(buf, ox, oy)
+
+
+class Tree(Diagram):
+    """Tidy tree of arbitrary nodes.
+
+    ``Tree(root, a, Tree(b, c, d), (e, f))`` — children are nodes, nested Trees or
+    ``(node, *children)`` tuples. direction: down | up | right | left. gap: space
+    between neighbouring subtrees (contour to contour); level_gap: space between depth
+    levels. edges: "ortho" (one Fork bus per parent), "straight" (outline-to-outline
+    lines), "none". edge_style: Edge / Fork keyword dict (colour, width, style, head …;
+    trees default to no arrowheads). ``tree.style_edge(child, label=…, style="dashed")``
+    gives the edge into one child its own style / label. align: "center" (parent over the
+    middle of its children) or "first" (over the first child).
+    """
+
+    def __init__(self, root: Node, *children: Any, direction: str = "down",
+                 gap: Optional[float] = None, level_gap: Optional[float] = None,
+                 edges: str = "ortho", edge_style: Optional[Dict[str, Any]] = None,
+                 align: str = "center", **kwargs: Any) -> None:
+        if direction not in _DIRS:
+            raise ValueError(f"direction must be one of {_DIRS}")
+        if edges not in ("ortho", "straight", "none"):
+            raise ValueError("edges must be ortho|straight|none")
+        if align not in ("center", "first"):
+            raise ValueError("align must be center|first")
+        self._spec = _Spec(root, [_as_spec(c) for c in children])
+        self.direction = direction
+        self.align = align
+        self.edge_mode = edges
+        self.edge_style = dict(edge_style or {})
+        theme = kwargs.get("theme")
+        self._gap_raw = None if gap is None else float(gap)
+        self._level_gap_raw = None if level_gap is None else float(level_gap)
+        super().__init__(_TreeLayout(self, theme), [], **kwargs)
+        self._rebuild_edges()
+
+    @property
+    def gap(self) -> float:
+        """Space between neighbouring subtrees (default ``theme.gap``)."""
+        return self._gap_raw if self._gap_raw is not None else float(self.theme.gap)
+
+    @gap.setter
+    def gap(self, v: Optional[float]) -> None:
+        self._gap_raw = None if v is None else float(v)
+
+    @property
+    def level_gap(self) -> float:
+        """Space between depth levels (default ``2.2 × theme.gap``)."""
+        return (self._level_gap_raw if self._level_gap_raw is not None
+                else float(self.theme.gap) * 2.2)
+
+    @level_gap.setter
+    def level_gap(self, v: Optional[float]) -> None:
+        self._level_gap_raw = None if v is None else float(v)
+
+    # ---- structure -------------------------------------------------------------
+    def _walk(self, s: Optional[_Spec] = None, parent: Optional[_Spec] = None):
+        s = s or self._spec
+        yield s, parent
+        for k in s.kids:
+            yield from self._walk(k, s)
+
+    def nodes(self) -> List[Node]:
+        return [s.node for s, _p in self._walk()]
+
+    def parent_of(self, node: Node) -> Optional[Node]:
+        for s, p in self._walk():
+            if s.node is node:
+                return p.node if p is not None else None
+        raise KeyError("node is not in this tree")
+
+    def style_edge(self, child: Node, **kw: Any) -> "Tree":
+        """Own style / label for the edge into ``child`` (e.g. ``style="dashed", label="evicted"``)."""
+        for s, p in self._walk():
+            if s.node is child:
+                if p is None:
+                    raise ValueError("the root has no incoming edge")
+                s.edge_kw.update(kw)
+                self._rebuild_edges()
+                return self
+        raise KeyError("node is not in this tree")
+
+    def _sides(self) -> Tuple[str, str]:
+        return {"down": ("bottom", "top"), "up": ("top", "bottom"),
+                "right": ("right", "left"), "left": ("left", "right")}[self.direction]
+
+    def _rebuild_edges(self) -> None:
+        out: List[Any] = []
+        if self.edge_mode == "none":
+            self.edges = out
+            return
+        base = {"head": "none"}
+        base.update(self.edge_style)
+        ps, cs = self._sides()
+        for s, _p in self._walk():
+            if not s.kids:
+                continue
+            if self.edge_mode == "straight":
+                for k in s.kids:
+                    kw = dict(base, **k.edge_kw)
+                    out.append(Edge(s.node.port("auto"), k.node.port("auto"), route="straight",
+                                    stub=0, **kw))
+                continue
+            labels = [k.edge_kw.get("label") for k in s.kids]
+            styles = [{x: v for x, v in k.edge_kw.items() if x != "label"} or None
+                      for k in s.kids]
+            for st_, lab in zip(styles, labels):  # branch labels sit beside the branch
+                if lab is not None and "label_side" not in st_:
+                    st_["label_side"] = "auto"
+            out.append(Fork(s.node.port(ps), [k.node.port(cs) for k in s.kids],
+                            labels=labels if any(l is not None for l in labels) else None,
+                            leaf_styles=styles if any(styles) else None, **base))
+        self.edges = out
+
+
+__all__ = ["Tree"]

codefig/venues.py

@@ -0,0 +1,93 @@
+"""Venue presets: final printed widths (inches) and figure rules of common templates.
+
+Numbers come from the official LaTeX templates (see ``docs/concepts/themes-venues.md``
+and the research notes ``figure_style_rules.md`` §4). Widths are in inches; codefig
+lays out in points (1 in = 72 pt), so ``Figure(venue="icml", width="column")`` is
+234 pt wide and every font size you see in the theme is the printed size.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Tuple, Union
+
+PT_PER_IN = 72.0
+
+
+@dataclass(frozen=True)
+class Venue:
+    """A paper template: text / column width (inches) and figure constraints."""
+
+    name: str
+    text: float
+    column: Optional[float] = None
+    body_pt: float = 10.0
+    font_rule: str = "type1|truetype"
+    min_font_pt: float = 6.0
+    min_stroke_pt: float = 0.5
+    strict_stroke: bool = False
+    no_title: bool = False
+    no_gray_text_bg: bool = False
+    bw_legible: bool = True
+    alt_text: bool = False
+    aliases: Tuple[str, ...] = field(default=())
+
+    def width_in(self, width: Union[str, float, None] = "text") -> float:
+        """Width in inches for ``"text"`` / ``"column"`` / a fraction ≤ 1 of the
+        text width / an absolute number of inches (> 1)."""
+        if width is None or width == "text":
+            return self.text
+        if width == "column":
+            if self.column is None:
+                raise ValueError(f"venue {self.name!r} is single-column; use width='text' "
+                                 "or a fraction such as width=0.5")
+            return self.column
+        w = float(width)
+        if w <= 0:
+            raise ValueError("width must be positive")
+        return self.text * w if w <= 1.0 else w
+
+    def width_pt(self, width: Union[str, float, None] = "text") -> float:
+        return self.width_in(width) * PT_PER_IN
+
+
+VENUES: Dict[str, Venue] = {
+    "neurips": Venue("neurips", 5.5, None, 10),
+    "iclr": Venue("iclr", 5.5, None, 10),
+    "icml": Venue("icml", 6.75, 3.25, 10, font_rule="type1", strict_stroke=True,
+                  no_title=True, no_gray_text_bg=True),
+    "mlsys": Venue("mlsys", 6.75, 3.25, 10, font_rule="type1", strict_stroke=True,
+                   no_title=True, no_gray_text_bg=True),
+    "vldb": Venue("vldb", 7.0, 3.33, 9, alt_text=True),
+    "sigmod": Venue("sigmod", 7.0, 3.33, 9, alt_text=True, aliases=("acm",)),
+    "osdi": Venue("osdi", 7.0, 3.335, 10, aliases=("usenix", "nsdi", "fast", "atc")),
+    "sosp": Venue("sosp", 7.0, 3.33, 10, alt_text=True),
+}
+_ALIASES = {a: v.name for v in VENUES.values() for a in v.aliases}
+
+
+def get_venue(venue: Union[str, Venue, None]) -> Optional[Venue]:
+    """``"icml"`` / ``"acm"`` / ``"usenix"`` / a :class:`Venue` → Venue (None → None)."""
+    if venue is None or isinstance(venue, Venue):
+        return venue
+    key = str(venue).lower()
+    key = _ALIASES.get(key, key)
+    if key not in VENUES:
+        names = sorted(set(VENUES) | set(_ALIASES))
+        raise ValueError(f"unknown venue {venue!r}; choose one of {names}")
+    return VENUES[key]
+
+
+def resolve_width_pt(venue: Union[str, Venue, None], width: Union[str, float, None]) -> Optional[float]:
+    """Final figure width in pt from ``venue`` + ``width`` (None if neither is given).
+
+    Without a venue, ``width`` must be a number of inches.
+    """
+    v = get_venue(venue)
+    if v is not None:
+        return v.width_pt(width)
+    if width is None:
+        return None
+    if isinstance(width, str):
+        raise ValueError(f"width={width!r} needs venue=; or pass inches, e.g. width=3.3")
+    return float(width) * PT_PER_IN