plot3 0.4.0__py3-none-any.whl

plot3/masking.py

@@ -0,0 +1,494 @@
+"""R-style bare-name / backtick column masking for plot3 (Jupyter layer).
+
+Mirrors tidy3 so ggplot2-style aesthetics work without string quotes::
+
+    aes(x=wt, y=mpg, colour=cyl)
+    aes(x=`First Name`, y=`Age (%)`)
+    facet_wrap(cyl)
+
+Two-phase design (SolveIt / IPython)::
+
+1. **Source preparser** turns backticks into a sentinel::
+
+       `First Name`  →  __plot3_bt__("First Name")
+
+2. **AST transformer** resolves the sentinel (and bare names) inside
+   :func:`aes` / :func:`facet_wrap` as **column-name strings**::
+
+       aes(x=wt, y=mpg)              → aes(x="wt", y="mpg")
+       aes(x=`First Name`, y=mpg)    → aes(x="First Name", y="mpg")
+       facet_wrap(cyl)               → facet_wrap("cyl")
+
+Plain ``.py`` files are unchanged — keep ``aes(x="wt", y="mpg")`` there.
+
+When tidy3 is also loaded, its backtick preparser may emit
+``__tidy3_bt__("name")``; this transformer treats that sentinel the same way
+inside ``aes`` / ``facet_wrap`` so dual-stack notebooks keep working.
+"""
+
+from __future__ import annotations
+
+import ast
+import builtins
+import io
+import re
+import tokenize
+from typing import Any, Iterable
+
+# Sentinels (plot3 own + tidy3 compatibility).
+BT_NAME = "__plot3_bt__"
+TIDY3_BT_NAME = "__tidy3_bt__"
+_BT_SENTINELS = frozenset({BT_NAME, TIDY3_BT_NAME})
+
+# Calls whose arguments / keywords are column *selectors* (→ strings).
+_SELECTOR_FUNCS = frozenset(
+    {
+        "aes",
+        "vars",
+        "facet_wrap",
+        "facet_grid",
+        "transition_time",
+        "transition_states",
+        "slider",
+    }
+)
+
+# Keywords that are plain labels / options — never rewrite as column names.
+_PASSTHROUGH_KW = frozenset(
+    {
+        "title",
+        "subtitle",
+        "caption",
+        "label",
+        "labels",
+        "trans",
+        "limits",
+        "palette",
+        "option",
+        "scales",
+        "ncol",
+        "nrow",
+        "bins",
+        "binwidth",
+        "width",
+        "size",
+        "alpha",
+        "color",  # const color on geoms is a colour code, not a column
+        "colour",
+        "linewidth",
+        "wireframe",
+        "levels",
+        "n",
+        "bw",
+        "kernel",
+        "trim",
+        "coef",
+        "varwidth",
+        "outlier",
+        "na_rm",
+        "hide",
+        "height",
+        "theme",
+        "kind",
+        "max_points",
+        "size_mode",
+        "fov",
+        "near",
+        "far",
+        "target",
+        "up",
+        "position",
+        "mark",
+        "at",
+        "stream",
+        "tlim",
+        "baseline",
+    }
+)
+
+_FORMULA_CALLS = frozenset({"geom_function", "geom_vector_field"})
+
+_BT_RE = re.compile(r"`([^`\n]+)`")
+
+
+def rewrite_formula_carets(source: str) -> str:
+    """Inside ``geom_function(...)``, turn ``^`` into ``**`` before Python parses.
+
+    ``2*x^3`` and ``(2*x)^3`` are the same tree once Python has parsed ``^``
+    as xor, so the rewrite has to happen on the text.
+    """
+    if "^" not in source or not any(name in source for name in _FORMULA_CALLS):
+        return source
+    try:
+        tokens = list(tokenize.generate_tokens(io.StringIO(source).readline))
+    except (tokenize.TokenError, SyntaxError):
+        return source
+    starts = [0]
+    for line in source.splitlines(keepends=True):
+        starts.append(starts[-1] + len(line))
+    hits: list[int] = []
+    depth = 0  # paren depth inside the current geom_function call
+    armed = False  # just saw the name geom_function
+    for tok in tokens:
+        if depth == 0:
+            if tok.type == tokenize.NAME and tok.string in _FORMULA_CALLS:
+                armed = True
+                continue
+            if armed and tok.type == tokenize.OP and tok.string == "(":
+                depth = 1
+            armed = False
+            continue
+        if tok.type == tokenize.OP:
+            if tok.string in "([{":
+                depth += 1
+            elif tok.string in ")]}":
+                depth -= 1
+            elif tok.string == "^":
+                row, col = tok.start
+                hits.append(starts[row - 1] + col)
+    for pos in reversed(hits):
+        source = source[:pos] + "**" + source[pos + 1 :]
+    return source
+
+
+def rewrite_backticks(source: str) -> str:
+    """Preparse backticks, and ``^`` inside ``geom_function``, before parsing.
+
+    `` `col name` `` → ``__plot3_bt__("col name")``. ``^`` → ``**`` only
+    inside a ``geom_function(...)`` call, where it means power.
+    """
+    source = rewrite_formula_carets(source)
+    if "`" not in source:
+        return source
+    return _BT_RE.sub(lambda m: f"{BT_NAME}({m.group(1)!r})", source)
+
+
+def plot3_backtick_transform(lines: list[str]) -> list[str]:
+    """IPython input transformer: backticks, and ``^`` inside geom_function."""
+    if not lines:
+        return lines
+    src = "".join(lines)
+    if "`" not in src and "^" not in src:
+        return lines
+    out = rewrite_backticks(src)
+    if out == src:
+        return lines
+    if out.endswith("\n"):
+        return out.splitlines(keepends=True)
+    parts = [ln + "\n" for ln in out.splitlines()]
+    return parts or [out]
+
+
+def _is_bt_call(node: ast.AST) -> bool:
+    return (
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id in _BT_SENTINELS
+        and len(node.args) == 1
+        and isinstance(node.args[0], ast.Constant)
+        and isinstance(node.args[0].value, str)
+    )
+
+
+def _string_const(name: str, old: ast.AST) -> ast.AST:
+    return ast.copy_location(ast.Constant(value=name), old)
+
+
+class MaskSelectors(ast.NodeTransformer):
+    """Rewrite bare names / backtick sentinels to string constants."""
+
+    def __init__(self, known: set[str]):
+        self.known = known
+
+    def visit_Name(self, node: ast.Name) -> ast.AST:
+        if isinstance(node.ctx, ast.Load) and node.id not in self.known:
+            return _string_const(node.id, node)
+        return node
+
+    def visit_Call(self, node: ast.Call) -> ast.AST:
+        if _is_bt_call(node):
+            return _string_const(str(node.args[0].value), node)
+
+        # Nested calls (e.g. unlikely helpers): keep func name; mask args.
+        if not isinstance(node.func, ast.Name):
+            node.func = self.visit(node.func)
+        node.args = [self.visit(arg) for arg in node.args]
+        node.keywords = [
+            ast.keyword(arg=kw.arg, value=self.visit(kw.value))
+            for kw in node.keywords
+        ]
+        return node
+
+
+def default_known_names(extra: Iterable[str] | None = None) -> set[str]:
+    """Names that must not become column strings (funcs, builtins, API)."""
+    known = set(dir(builtins))
+    known.update(_BT_SENTINELS)
+    known.update({"True", "False", "None"})
+    try:
+        import plot3 as p3
+
+        for name in getattr(p3, "__all__", ()):
+            if not str(name).startswith("_"):
+                known.add(name)
+    except Exception:
+        pass
+    # Common frame / helper names left alone when present in the user ns
+    # are added via extra from IPython.user_ns.
+    if extra:
+        known.update(extra)
+    return known
+
+
+class Plot3MaskTransformer(ast.NodeTransformer):
+    """AST pass: bare names / backticks → column strings inside aes / facet_wrap."""
+
+    def __init__(self, known: set[str] | None = None):
+        self._known_static = known
+
+    def _known(self) -> set[str]:
+        if self._known_static is not None:
+            return set(self._known_static)
+        extra: set[str] = set()
+        try:
+            from IPython import get_ipython
+
+            ip = get_ipython()
+            if ip is not None and getattr(ip, "user_ns", None) is not None:
+                extra.update(ip.user_ns.keys())
+        except Exception:
+            pass
+        return default_known_names(extra)
+
+    def _mask_selector(self, node: ast.AST) -> ast.AST:
+        return MaskSelectors(self._known()).visit(node)
+
+    def visit_Call(self, node: ast.Call) -> ast.AST:
+        # Always recurse so nested aes(...) inside geom_point(aes(...)) is seen.
+        if not isinstance(node.func, ast.Name):
+            return self.generic_visit(node)
+
+        name = node.func.id
+        if name in _FORMULA_CALLS:
+            return self._visit_geom_function(node)
+        if name not in _SELECTOR_FUNCS:
+            return self.generic_visit(node)
+
+        if name == "aes":
+            # All positional + keyword values are column selectors.
+            node.args = [self._mask_selector(a) for a in node.args]
+            node.keywords = [
+                ast.keyword(arg=kw.arg, value=self._mask_selector(kw.value))
+                for kw in node.keywords
+            ]
+            return node
+
+        if name == "transition_time":
+            # The column is a selector. Parameter ranges (``a=(0, 2*pi)``)
+            # and ``frames=`` stay Python, so ``pi`` is not quoted.
+            node.args = [self._mask_selector(a) for a in node.args]
+            new_kws: list[ast.keyword] = []
+            for kw in node.keywords:
+                if kw.arg == "column":
+                    new_kws.append(
+                        ast.keyword(
+                            arg=kw.arg, value=self._mask_selector(kw.value)
+                        )
+                    )
+                else:
+                    new_kws.append(
+                        ast.keyword(arg=kw.arg, value=self.visit(kw.value))
+                    )
+            node.keywords = new_kws
+            return node
+
+        if name == "slider":
+            # Ranges and ``steps=`` stay Python, so ``pi`` is not quoted.
+            node.args = [self.visit(a) for a in node.args]
+            node.keywords = [
+                ast.keyword(arg=kw.arg, value=self.visit(kw.value))
+                for kw in node.keywords
+            ]
+            return node
+
+        if name == "transition_states":
+            # The frame column is a selector, positional or keyword.
+            node.args = [self._mask_selector(a) for a in node.args]
+            node.keywords = [
+                ast.keyword(arg=kw.arg, value=self._mask_selector(kw.value))
+                for kw in node.keywords
+            ]
+            return node
+
+        if name in {"facet_wrap", "facet_grid"}:
+            # facets / rows / cols are column selectors; ncol, nrow, and
+            # scales stay as-is.
+            node.args = [self._mask_selector(a) for a in node.args]
+            new_kws: list[ast.keyword] = []
+            for kw in node.keywords:
+                if kw.arg in (None, "facets") or (
+                    kw.arg is not None and kw.arg not in _PASSTHROUGH_KW
+                ):
+                    new_kws.append(
+                        ast.keyword(arg=kw.arg, value=self._mask_selector(kw.value))
+                    )
+                else:
+                    new_kws.append(kw)
+            node.keywords = new_kws
+            return node
+
+        return self.generic_visit(node)
+
+    def _visit_geom_function(self, node: ast.Call) -> ast.AST:
+        """Stringify formula arguments (``x=``, ``y=``, ``z=``, ``f=``, ``dx=``, ``dy=``).
+
+        A parametric curve and a vector field quote every formula keyword.
+        A number such as ``z=1`` stays Python, so it is a coefficient and
+        not a second formula. ``a=``, ``xlim=``, ``n=``, ``mark=``, and
+        ``tlim=`` stay Python too. ``^`` is already ``**`` if the text hook ran.
+        """
+        known = self._known()
+        call_name = node.func.id if isinstance(node.func, ast.Name) else ""
+        axis_names = (
+            {"dx", "dy", "dz"}
+            if call_name == "geom_vector_field"
+            else {"x", "y", "z", "f"}
+        )
+        formula_done = False
+        new_args: list[ast.AST] = []
+        for index, arg in enumerate(node.args):
+            if index == 0 and not formula_done and _is_formula_expr(arg, known):
+                new_args.append(_formula_string(arg))
+                formula_done = True
+            else:
+                new_args.append(self.visit(arg))
+        expr_axes: list[ast.keyword] = []
+        numeric_axes: list[ast.keyword] = []
+        for kw in node.keywords:
+            if kw.arg in axis_names and _is_number_const(kw.value):
+                numeric_axes.append(kw)
+            elif kw.arg in axis_names and _is_formula_expr(kw.value, known):
+                expr_axes.append(kw)
+        # geom_function(y=2) is a constant formula. z=1 next to a real
+        # formula is a coefficient, so it stays a number.
+        quote_numeric = (
+            call_name == "geom_function"
+            and not formula_done
+            and not expr_axes
+            and len(numeric_axes) == 1
+        )
+        quoted = {id(kw) for kw in expr_axes}
+        if quote_numeric:
+            quoted.add(id(numeric_axes[0]))
+        new_kws: list[ast.keyword] = []
+        for kw in node.keywords:
+            if kw.arg == "mark" and _mark_is_bare(kw.value):
+                new_kws.append(
+                    ast.keyword(arg="mark", value=_quote_mark(kw.value))
+                )
+            elif id(kw) in quoted:
+                new_kws.append(
+                    ast.keyword(arg=kw.arg, value=_formula_string(kw.value))
+                )
+            else:
+                new_kws.append(
+                    ast.keyword(arg=kw.arg, value=self.visit(kw.value))
+                )
+        node.args = new_args
+        node.keywords = new_kws
+        return node
+
+
+def _is_number_const(node: ast.AST) -> bool:
+    """True for ``1`` and ``-1``, which are coefficients rather than formulas."""
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)):
+        node = node.operand
+    return (
+        isinstance(node, ast.Constant)
+        and isinstance(node.value, (int, float))
+        and not isinstance(node.value, bool)
+    )
+
+
+def _mark_is_bare(node: ast.AST) -> bool:
+    if isinstance(node, ast.Name):
+        return True
+    if isinstance(node, (ast.List, ast.Tuple)):
+        return bool(node.elts) and all(isinstance(elt, ast.Name) for elt in node.elts)
+    return False
+
+
+def _quote_mark(node: ast.AST) -> ast.AST:
+    if isinstance(node, ast.Name):
+        return ast.copy_location(ast.Constant(value=node.id), node)
+    elts = [
+        ast.copy_location(ast.Constant(value=elt.id), elt)
+        if isinstance(elt, ast.Name)
+        else elt
+        for elt in node.elts
+    ]
+    copied = ast.List(elts=elts, ctx=ast.Load()) if isinstance(node, ast.List) else ast.Tuple(elts=elts, ctx=ast.Load())
+    return ast.copy_location(copied, node)
+
+
+def _is_formula_expr(node: ast.AST, known: set[str]) -> bool:
+    """True when this argument is a formula to quote, not a Python value."""
+    if isinstance(node, ast.Lambda):
+        return False
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return False
+    if isinstance(node, (ast.Attribute, ast.Subscript)):
+        return False
+    if isinstance(node, ast.Name) and node.id in known:
+        return False
+    if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+        return False
+    return True
+
+
+def _formula_string(node: ast.AST) -> ast.Constant:
+    # ``^`` left in the tree is xor, and ``2*x^3`` is already ``(2*x)^3``.
+    # Do not guess. The text hook should have rewritten it.
+    if any(isinstance(child, ast.BitXor) for child in ast.walk(node)):
+        raise ValueError('use ** or quote the formula: "y = x^2"')
+    text = ast.unparse(node)
+    return ast.copy_location(ast.Constant(value=text), node)
+
+
+def apply_masking(
+    source: str,
+    *,
+    known: set[str] | None = None,
+    backticks: bool = True,
+) -> str:
+    """Apply backtick rewrite + AST masking; return unparsed source (tests)."""
+    text = rewrite_formula_carets(source)
+    if backticks:
+        text = rewrite_backticks(text)
+    tree = ast.parse(text)
+    tree = Plot3MaskTransformer(known=known or default_known_names()).visit(tree)
+    ast.fix_missing_locations(tree)
+    return ast.unparse(tree)
+
+
+def is_mask_transformer(obj: Any) -> bool:
+    return isinstance(obj, Plot3MaskTransformer) or (
+        type(obj).__name__ == "Plot3MaskTransformer"
+        and getattr(type(obj), "__module__", "").startswith("plot3")
+    )
+
+
+def is_backtick_transformer(obj: Any) -> bool:
+    return (
+        getattr(obj, "__module__", "") in {"plot3.masking", "plot3.jupyter"}
+        and getattr(obj, "__name__", "") == "plot3_backtick_transform"
+    )
+
+
+def is_tidy3_backtick_transformer(obj: Any) -> bool:
+    """True if tidy3 already installed a backtick preparser (share it)."""
+    mod = getattr(obj, "__module__", "") or ""
+    name = getattr(obj, "__name__", "") or ""
+    return name.endswith("backtick_transform") and (
+        mod.startswith("tidy3") or name == "tidy3_backtick_transform"
+    )