redis-lua-py 0.1.0__py3-none-any.whl

redis_lua_py/_lua.py

@@ -0,0 +1,308 @@
+"""A small Lua IR and the emitter that turns it into source text.
+
+The IR is deliberately Lua-shaped rather than Python-shaped: every semantic
+gap between the two languages is closed in ``_compile``, so that by the time a
+node reaches this module it means exactly what the emitted text means. That
+keeps the emitter dumb, which is what makes the output reviewable.
+
+This module is private. Its shape is not part of the public API.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+# Lua 5.1 operator precedence, lowest binding first. Used to decide where
+# parentheses are actually required, so the emitted source stays readable.
+_PRECEDENCE: dict[str, int] = {
+    "or": 1,
+    "and": 2,
+    "<": 3,
+    ">": 3,
+    "<=": 3,
+    ">=": 3,
+    "~=": 3,
+    "==": 3,
+    "..": 4,
+    "+": 5,
+    "-": 5,
+    "*": 6,
+    "/": 6,
+    "%": 6,
+    "^": 8,
+}
+_UNARY_PRECEDENCE = 7
+_RIGHT_ASSOCIATIVE = {"..", "^"}
+
+_LUA_KEYWORDS = frozenset(
+    [
+        "and",
+        "break",
+        "do",
+        "else",
+        "elseif",
+        "end",
+        "false",
+        "for",
+        "function",
+        "if",
+        "in",
+        "local",
+        "nil",
+        "not",
+        "or",
+        "repeat",
+        "return",
+        "then",
+        "true",
+        "until",
+        "while",
+    ]
+)
+
+_ESCAPES = {"\\": "\\\\", "'": "\\'", "\n": "\\n", "\r": "\\r", "\t": "\\t", "\0": "\\0"}
+
+
+def quote(value: str) -> str:
+    """Render a Python string as a single-quoted Lua literal."""
+    out = []
+    for ch in value:
+        if ch in _ESCAPES:
+            out.append(_ESCAPES[ch])
+        elif ord(ch) < 0x20 or ord(ch) == 0x7F:
+            out.append(f"\\{ord(ch):03d}")
+        else:
+            out.append(ch)
+    return "'" + "".join(out) + "'"
+
+
+def is_identifier(name: str) -> bool:
+    return name.isidentifier() and name not in _LUA_KEYWORDS
+
+
+class Node:
+    __slots__ = ()
+
+
+class Expr(Node):
+    __slots__ = ()
+
+
+class Stat(Node):
+    __slots__ = ()
+
+
+@dataclass(frozen=True, slots=True)
+class Nil(Expr):
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class Bool(Expr):
+    value: bool
+
+
+@dataclass(frozen=True, slots=True)
+class Num(Expr):
+    value: int | float
+
+
+@dataclass(frozen=True, slots=True)
+class Str(Expr):
+    value: str
+
+
+@dataclass(frozen=True, slots=True)
+class Name(Expr):
+    id: str
+
+
+@dataclass(frozen=True, slots=True)
+class BinOp(Expr):
+    op: str
+    left: Expr
+    right: Expr
+
+
+@dataclass(frozen=True, slots=True)
+class UnOp(Expr):
+    op: str  # "not", "-", "#"
+    operand: Expr
+
+
+@dataclass(frozen=True, slots=True)
+class Index(Expr):
+    obj: Expr
+    key: Expr
+
+
+@dataclass(frozen=True, slots=True)
+class Call(Expr):
+    func: Expr
+    args: tuple[Expr, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class Table(Expr):
+    array: tuple[Expr, ...] = ()
+    hash: tuple[tuple[Expr, Expr], ...] = ()
+
+
+@dataclass(slots=True)
+class Local(Stat):
+    names: list[str]
+    values: list[Expr]
+
+
+@dataclass(slots=True)
+class Assign(Stat):
+    targets: list[Expr]
+    values: list[Expr]
+
+
+@dataclass(slots=True)
+class If(Stat):
+    branches: list[tuple[Expr, list[Stat]]]
+    orelse: list[Stat] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class NumericFor(Stat):
+    var: str
+    start: Expr
+    stop: Expr
+    step: Expr | None
+    body: list[Stat]
+
+
+@dataclass(slots=True)
+class While(Stat):
+    test: Expr
+    body: list[Stat]
+
+
+@dataclass(slots=True)
+class Return(Stat):
+    value: Expr | None = None
+
+
+@dataclass(slots=True)
+class Break(Stat):
+    pass
+
+
+@dataclass(slots=True)
+class ExprStat(Stat):
+    expr: Expr
+
+
+@dataclass(slots=True)
+class Comment(Stat):
+    text: str
+
+
+def emit_expr(node: Expr, parent_prec: int = 0) -> str:
+    """Render an expression, parenthesising only where precedence demands it."""
+    match node:
+        case Nil():
+            return "nil"
+        case Bool(value=v):
+            return "true" if v else "false"
+        case Num(value=v):
+            return repr(v) if isinstance(v, float) else str(v)
+        case Str(value=v):
+            return quote(v)
+        case Name(id=v):
+            return v
+        case Table(array=arr, hash=pairs):
+            items = [emit_expr(a) for a in arr]
+            items += [
+                (
+                    f"[{emit_expr(k)}] = {emit_expr(val)}"
+                    if not (isinstance(k, Str) and is_identifier(k.value))
+                    else f"{k.value} = {emit_expr(val)}"
+                )
+                for k, val in pairs
+            ]
+            return "{" + ", ".join(items) + "}"
+        case Index(obj=obj, key=key):
+            base = emit_expr(obj, 9)
+            if isinstance(key, Str) and is_identifier(key.value):
+                return f"{base}.{key.value}"
+            return f"{base}[{emit_expr(key)}]"
+        case Call(func=func, args=args):
+            rendered = ", ".join(emit_expr(a) for a in args)
+            return f"{emit_expr(func, 9)}({rendered})"
+        case UnOp(op=op, operand=operand):
+            spacer = " " if op == "not" else ""
+            text = f"{op}{spacer}{emit_expr(operand, _UNARY_PRECEDENCE)}"
+            return f"({text})" if parent_prec > _UNARY_PRECEDENCE else text
+        case BinOp(op=op, left=left, right=right):
+            prec = _PRECEDENCE[op]
+            # For a right-associative operator the left operand needs the
+            # tighter bound, and vice versa.
+            if op in _RIGHT_ASSOCIATIVE:
+                left_prec, right_prec = prec + 1, prec
+            else:
+                left_prec, right_prec = prec, prec + 1
+            text = f"{emit_expr(left, left_prec)} {op} {emit_expr(right, right_prec)}"
+            return f"({text})" if prec < parent_prec else text
+        case _:  # pragma: no cover - guards against an IR node with no emitter
+            raise TypeError(f"cannot emit expression node {type(node).__name__}")
+
+
+def emit_block(body: list[Stat], indent: int = 0) -> list[str]:
+    lines: list[str] = []
+    pad = "  " * indent
+    for stat in body:
+        match stat:
+            case Comment(text=text):
+                lines += [f"{pad}-- {line}" for line in text.splitlines()]
+            case Local(names=names, values=values):
+                target = ", ".join(names)
+                if values:
+                    rhs = ", ".join(emit_expr(v) for v in values)
+                    lines.append(f"{pad}local {target} = {rhs}")
+                else:
+                    lines.append(f"{pad}local {target}")
+            case Assign(targets=targets, values=values):
+                lhs = ", ".join(emit_expr(t) for t in targets)
+                rhs = ", ".join(emit_expr(v) for v in values)
+                lines.append(f"{pad}{lhs} = {rhs}")
+            case ExprStat(expr=expr):
+                lines.append(f"{pad}{emit_expr(expr)}")
+            case Return(value=value):
+                rendered = "" if value is None else f" {emit_expr(value)}"
+                lines.append(f"{pad}return{rendered}")
+            case Break():
+                # Lua 5.1 requires `break` to be the final statement of a
+                # block; wrapping it in `do ... end` makes that true wherever
+                # Python allowed it.
+                lines.append(f"{pad}do break end")
+            case While(test=test, body=inner):
+                lines.append(f"{pad}while {emit_expr(test)} do")
+                lines += emit_block(inner, indent + 1)
+                lines.append(f"{pad}end")
+            case NumericFor(var=var, start=start, stop=stop, step=step, body=inner):
+                header = f"{pad}for {var} = {emit_expr(start)}, {emit_expr(stop)}"
+                if step is not None:
+                    header += f", {emit_expr(step)}"
+                lines.append(header + " do")
+                lines += emit_block(inner, indent + 1)
+                lines.append(f"{pad}end")
+            case If(branches=branches, orelse=orelse):
+                for i, (test, inner) in enumerate(branches):
+                    keyword = "if" if i == 0 else "elseif"
+                    lines.append(f"{pad}{keyword} {emit_expr(test)} then")
+                    lines += emit_block(inner, indent + 1)
+                if orelse:
+                    lines.append(f"{pad}else")
+                    lines += emit_block(orelse, indent + 1)
+                lines.append(f"{pad}end")
+            case _:  # pragma: no cover - guards against an IR node with no emitter
+                raise TypeError(f"cannot emit statement node {type(stat).__name__}")
+    return lines
+
+
+def emit(body: list[Stat]) -> str:
+    return "\n".join(emit_block(body)) + "\n"

redis_lua_py/_runtime.py

@@ -0,0 +1,72 @@
+"""Names that script bodies refer to.
+
+Nothing here ever runs. A decorated function's body is read as source and
+compiled to Lua; Python never executes it. These objects exist so that the
+body is a valid Python expression to your editor, your linter and your type
+checker, and so that calling one by mistake fails loudly instead of silently
+doing nothing.
+"""
+
+from __future__ import annotations
+
+from typing import Any, NoReturn
+
+
+class Key(str):
+    """Marks a parameter as a Redis key.
+
+    Parameters annotated ``Key`` become ``KEYS[n]`` inside the script, in
+    declaration order; everything else becomes ``ARGV[n]``. Getting this right
+    matters: Redis Cluster routes a script by its declared keys, so a key
+    passed as an argument will be invisible to the router.
+    """
+
+    __slots__ = ()
+
+
+class _Namespace:
+    """Attribute access that type checkers accept and runtime refuses.
+
+    The compiler identifies these by value rather than by the name they are
+    imported under, so every alias works and nothing is reserved.
+    """
+
+    __slots__ = ("_name",)
+
+    def __init__(self, name: str) -> None:
+        object.__setattr__(self, "_name", name)
+
+    @property
+    def kind(self) -> str:
+        """Which Lua table this stands for: ``redis`` or ``cjson``."""
+        return str(object.__getattribute__(self, "_name"))
+
+    def __getattr__(self, item: str) -> Any:
+        namespace = object.__getattribute__(self, "_name")
+
+        def _stub(*_args: object, **_kwargs: object) -> NoReturn:
+            raise RuntimeError(
+                f"{namespace}.{item}() is a Lua construct and cannot run in Python. "
+                "It is only meaningful inside the body of an @script function, "
+                "which is compiled rather than executed."
+            )
+
+        return _stub
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        return f"<redis_lua_py {object.__getattribute__(self, '_name')} namespace>"
+
+
+#: Call Redis commands: ``redis.incr(key)`` becomes ``redis.call('INCR', key)``.
+#: Underscores split into subcommand tokens, so ``redis.script_load(x)``
+#: becomes ``redis.call('SCRIPT', 'LOAD', x)``.
+#:
+#: Import it under any name you like. If the module also imports the redis-py
+#: client, ``from redis_lua_py import redis as r`` keeps the two apart.
+redis = _Namespace("redis")
+
+#: Alias for :data:`redis`, for modules that would rather not rename anything.
+call = _Namespace("redis")
+
+#: The JSON library Redis exposes to scripts: ``cjson.encode`` / ``cjson.decode``.
+cjson = _Namespace("cjson")

redis_lua_py/_script.py

@@ -0,0 +1,151 @@
+"""The callable object a decorated function becomes."""
+
+from __future__ import annotations
+
+import difflib
+from dataclasses import dataclass, field
+from typing import Any
+from weakref import WeakKeyDictionary
+
+from .errors import ScriptArgumentError
+
+
+def encode(name: str, value: object) -> str | bytes | memoryview:
+    """Render a Python value as a Redis argument.
+
+    Redis has no argument types: everything on the wire is a byte string. This
+    only accepts values whose string form is unambiguous, so that a stray None
+    or object fails here rather than arriving in Lua as something surprising.
+    """
+    if isinstance(value, bool):
+        return "1" if value else "0"
+    if isinstance(value, str | bytes | memoryview):
+        return value
+    if isinstance(value, int | float):
+        return repr(value) if isinstance(value, float) else str(value)
+    raise ScriptArgumentError(
+        f"argument {name!r} is a {type(value).__name__}, which has no Redis representation. "
+        "Pass a str, bytes, int, float or bool."
+    )
+
+
+@dataclass(frozen=True)
+class CompiledScript:
+    """A Python function compiled to Lua, callable against a Redis client.
+
+    Calling it runs EVALSHA and falls back to EVAL the first time, or whenever
+    the server has dropped the script from its cache. Pass a sync client and
+    you get a value; pass an async client and you get an awaitable.
+    """
+
+    name: str
+    lua: str
+    params: tuple[str, ...]
+    keys: tuple[str, ...]
+    args: tuple[str, ...]
+    doc: str | None = None
+    source: str = ""
+    _registry: WeakKeyDictionary[Any, Any] = field(
+        default_factory=WeakKeyDictionary, compare=False, repr=False
+    )
+
+    def __call__(self, client: Any, /, *positional: object, **keyword: object) -> Any:
+        keys, argv = self.resolve(*positional, **keyword)
+        return self._for(client)(keys=keys, args=argv, client=client)
+
+    def bind(self, client: Any) -> BoundScript:
+        """Attach a client, so that callers stop repeating it.
+
+        The unbound form keeps working; this only removes the first argument.
+        Sync and async clients bind the same way, and the bound call returns
+        whatever the underlying client would have.
+        """
+        return BoundScript(self, client)
+
+    def resolve(self, *positional: object, **keyword: object) -> tuple[list[Any], list[Any]]:
+        """Resolve call arguments into the KEYS and ARGV lists."""
+        if len(positional) > len(self.params):
+            raise ScriptArgumentError(
+                f"{self.name}() takes {len(self.params)} argument(s), got {len(positional)}"
+            )
+
+        values: dict[str, object] = dict(zip(self.params, positional, strict=False))
+        for key, value in keyword.items():
+            if key not in self.params:
+                suggestion = difflib.get_close_matches(key, self.params, n=1)
+                hint = f"; did you mean {suggestion[0]!r}?" if suggestion else ""
+                raise ScriptArgumentError(
+                    f"{self.name}() has no parameter {key!r}{hint} "
+                    f"(parameters: {', '.join(self.params)})"
+                )
+            if key in values:
+                raise ScriptArgumentError(f"{self.name}() got two values for {key!r}")
+            values[key] = value
+
+        missing = [p for p in self.params if p not in values]
+        if missing:
+            raise ScriptArgumentError(f"{self.name}() is missing argument(s): {', '.join(missing)}")
+
+        return (
+            [values[k] for k in self.keys],
+            [encode(a, values[a]) for a in self.args],
+        )
+
+    def _for(self, client: Any) -> Any:
+        """Get the redis-py Script bound to this client, registering it once.
+
+        redis-py's own Script object already implements the EVALSHA-then-EVAL
+        dance and the NOSCRIPT retry, so this defers to it rather than
+        reimplementing script caching.
+        """
+        try:
+            registered = self._registry.get(client)
+        except TypeError:  # a client that does not support weak references
+            return client.register_script(self.lua)
+        if registered is None:
+            registered = client.register_script(self.lua)
+            self._registry[client] = registered
+        return registered
+
+    def __repr__(self) -> str:
+        signature = ", ".join(self.params)
+        return f"<script {self.name}({signature}) from {self.source}>"
+
+
+@dataclass(frozen=True)
+class BoundScript:
+    """A script with its client already attached, produced by :meth:`bind`."""
+
+    script: CompiledScript
+    client: Any
+
+    def __call__(self, *positional: object, **keyword: object) -> Any:
+        return self.script(self.client, *positional, **keyword)
+
+    @property
+    def name(self) -> str:
+        return self.script.name
+
+    @property
+    def lua(self) -> str:
+        return self.script.lua
+
+    @property
+    def params(self) -> tuple[str, ...]:
+        return self.script.params
+
+    @property
+    def keys(self) -> tuple[str, ...]:
+        return self.script.keys
+
+    @property
+    def args(self) -> tuple[str, ...]:
+        return self.script.args
+
+    @property
+    def doc(self) -> str | None:
+        return self.script.doc
+
+    def __repr__(self) -> str:
+        signature = ", ".join(self.script.params)
+        return f"<bound script {self.script.name}({signature})>"

redis_lua_py/errors.py

@@ -0,0 +1,53 @@
+"""Errors raised while compiling or running a script."""
+
+from __future__ import annotations
+
+
+class RedisLuaError(Exception):
+    """Base class for every error this package raises."""
+
+
+class CompileError(RedisLuaError):
+    """A Python function could not be turned into Lua."""
+
+
+class UnsupportedSyntax(CompileError):
+    """The function used Python that has no meaning inside a Redis script.
+
+    Carries the location so the message points at the offending line rather
+    than at the decorator.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        filename: str,
+        lineno: int,
+        col: int,
+        source_line: str | None = None,
+        hint: str | None = None,
+    ) -> None:
+        self.message = message
+        self.filename = filename
+        self.lineno = lineno
+        self.col = col
+        self.source_line = source_line
+        self.hint = hint
+        super().__init__(self._render())
+
+    def _render(self) -> str:
+        parts = [f"{self.message}", f'  File "{self.filename}", line {self.lineno}']
+        if self.source_line is not None:
+            stripped = self.source_line.rstrip()
+            indent = len(stripped) - len(stripped.lstrip())
+            parts.append(f"    {stripped.lstrip()}")
+            caret_col = max(self.col - indent, 0)
+            parts.append("    " + " " * caret_col + "^")
+        if self.hint:
+            parts.append(f"  hint: {self.hint}")
+        return "\n".join(parts)
+
+
+class ScriptArgumentError(RedisLuaError):
+    """A script was called with the wrong keys or arguments."""