opencode-pyneruntime 6.6.4__py3-none-any.whl

pynecore/transformers/module_property.py

@@ -0,0 +1,221 @@
+from typing import cast, Any
+import ast
+import json
+from pathlib import Path
+
+
+class ModulePropertyTransformer(ast.NodeTransformer):
+    """
+    Transform lib.xxx references based on the generated module_properties.json registry.
+
+    - ``property`` entries (Pine names that are values): bare reads become calls
+      (``ta.tr`` -> ``ta.tr()``); explicit calls are left untouched.
+    - ``variable`` entries: left as plain attribute reads.
+    - Function-and-namespace modules (``plot``, ``dayofweek``, ...): calls and promoted
+      bare reads are routed to the module's self-named function
+      (``plot(x)`` -> ``plot.plot(x)``, bare ``dayofweek`` -> ``dayofweek.dayofweek()``).
+    - Unknown names on known pynecore.lib modules raise at transform time — the
+      registry is exhaustive, so this catches typos and a stale registry early.
+    - Unknown module paths (user ``lib.*`` workdir libraries) and ``_``-prefixed
+      names are plain attribute reads.
+    """
+
+    def __init__(self):
+        # Structure: module -> name -> {"type": "property"|"variable"}
+        self.module_info: dict[str, dict[str, dict[str, Any]]] = {}
+
+        # Load config
+        try:
+            with open(Path(__file__).parent / "module_properties.json") as f:
+                self.module_info = json.load(f)
+        except (IOError, json.JSONDecodeError) as e:
+            raise RuntimeError(f"Failed to load module properties config: {e}")
+
+    def visit(self, node: ast.AST) -> ast.AST:
+        """
+        Override the generic visit method to set .parent on each child node
+        for chain detection.
+        """
+        # Set parent on children
+        for field, value in ast.iter_fields(node):
+            if isinstance(value, ast.AST):
+                setattr(value, "parent", node)
+            elif isinstance(value, list):
+                for item in value:
+                    if isinstance(item, ast.AST):
+                        setattr(item, "parent", node)
+
+        return super().visit(node)
+
+    def visit_Attribute(self, node: ast.Attribute) -> ast.AST:
+        """Process attribute access, but skip if inside type annotations."""
+        node = cast(ast.Attribute, self.generic_visit(node))
+
+        # Skip if inside type annotations
+        if self._is_in_type_annotation(node):
+            return node
+
+        # Retrieve the AST parent node
+        parent = getattr(node, 'parent', None)
+
+        # Intermediate module - if the parent is also an Attribute, this is not the topmost attribute
+        if isinstance(parent, ast.Attribute):
+            return node
+
+        # If this node has already been processed, or the chain does not start with lib..., skip
+        if hasattr(node, '_processed') or not self._is_lib_reference(node):
+            return node
+
+        # Now it's the topmost attribute (e.g., ...data_window)
+        # Check the full module path and the final attribute
+        module_path, name = self._get_module_info(node)
+        if not module_path or not name:
+            return node
+
+        full_path = f"{module_path}.{name}"
+
+        # Call site — explicit calls stay as they are, except when the callee is a
+        # function-and-namespace module (its registry entry contains a self-named
+        # function): ``lib.plot(...)`` routes to ``lib.plot.plot(...)``
+        if isinstance(parent, ast.Call) and parent.func == node:
+            inner_attrs = self.module_info.get(full_path)
+            if inner_attrs is not None and name in inner_attrs:
+                result: ast.expr = ast.Attribute(value=node, attr=name, ctx=ast.Load())
+                setattr(result, "_processed", True)
+                return result
+            return node
+
+        module_attrs = self.module_info.get(module_path)
+        if module_attrs is None:
+            # Unknown module path: a user workdir library or a class-attribute chain —
+            # plain Python attribute access
+            return node
+
+        attr_info = module_attrs.get(name)
+        if attr_info is not None:
+            if attr_info["type"] == "property":
+                if full_path == "lib.na":
+                    # Bare ``na`` is a constant value (the interned typeless NA):
+                    # load it directly instead of emitting a per-bar ``lib.na()``
+                    # call. Explicit ``na(x)`` predicate calls are untouched above.
+                    result = ast.Attribute(value=ast.Name(id='lib', ctx=ast.Load()),
+                                           attr='_na_none', ctx=ast.Load())
+                    setattr(result, "_processed", True)
+                    return result
+                inner_attrs = self.module_info.get(full_path)
+                if inner_attrs is not None and name in inner_attrs:
+                    # Promoted self-named property of a function-and-namespace module:
+                    # bare ``dayofweek`` -> ``lib.dayofweek.dayofweek()``
+                    func: ast.expr = ast.Attribute(value=self._copy_node(node), attr=name,
+                                                   ctx=ast.Load())
+                else:
+                    func = self._copy_node(node)
+                result = ast.Call(func=func, args=[], keywords=[])
+            else:
+                result = node
+
+        # Submodule reference — leave as is
+        elif full_path in self.module_info:
+            result = node
+
+        # Internal names are never module properties — plain attribute access
+        elif name.startswith('_'):
+            result = node
+
+        else:
+            raise SyntaxError(
+                f"unknown attribute '{name}' on module '{module_path}' (line {node.lineno}); "
+                f"if this is a new pynecore.lib name, regenerate module_properties.json "
+                f"with scripts/module_property_collector.py"
+            )
+
+        setattr(result, "_processed", True)
+        return result
+
+    @staticmethod
+    def _is_lib_reference(node: ast.Attribute) -> bool:
+        """Check if the attribute chain starts with 'lib'."""
+        current = node
+        while isinstance(current, ast.Attribute):
+            current = current.value
+        return isinstance(current, ast.Name) and current.id == 'lib'
+
+    @staticmethod
+    def _get_module_info(node: ast.Attribute) -> tuple[str | None, str | None]:
+        """
+        Gather the full chain of attributes until we reach 'lib',
+        then split into (module_path, final_attribute).
+        Example: lib.display.data_window -> (lib.display, data_window)
+        """
+        attrs = []
+        current = node
+        while isinstance(current, ast.Attribute):
+            attrs.append(current.attr)
+            current = current.value
+
+        if isinstance(current, ast.Name) and current.id == 'lib':
+            attrs.append('lib')
+            attrs.reverse()
+            # Example: ['lib', 'display', 'data_window']
+            if len(attrs) < 2:
+                return None, None
+            module_path = '.'.join(attrs[:-1])  # 'lib.display'
+            final_attr = attrs[-1]  # 'data_window'
+            return module_path, final_attr
+        return None, None
+
+    def _is_in_type_annotation(self, node: ast.Attribute) -> bool:
+        """Check if the node is inside a type annotation."""
+        current = node
+        while hasattr(current, 'parent'):
+            parent = getattr(current, 'parent', None)
+
+            # Check if we're in an annotated assignment's annotation
+            if (isinstance(parent, ast.AnnAssign) and parent.annotation and
+                    self._is_node_in_subtree(node, cast(ast.AST, parent.annotation))):
+                return True
+
+            # Check if we're in a function argument's annotation
+            if (isinstance(parent, ast.arg) and parent.annotation and
+                    self._is_node_in_subtree(node, cast(ast.AST, parent.annotation))):
+                return True
+
+            # Check if we're in a function return annotation
+            if (isinstance(parent, ast.FunctionDef) and parent.returns and
+                    self._is_node_in_subtree(node, parent.returns)):
+                return True
+
+            # Check if we're in an async function return annotation
+            if (isinstance(parent, ast.AsyncFunctionDef) and parent.returns and
+                    self._is_node_in_subtree(node, parent.returns)):
+                return True
+
+            current = parent
+
+        return False
+
+    @staticmethod
+    def _is_node_in_subtree(node: ast.AST, subtree: ast.AST | None) -> bool:
+        """Check if a node is contained within a subtree."""
+        if subtree is None:
+            return False
+
+        if node is subtree:
+            return True
+
+        # Recursively check all child nodes
+        for child in ast.walk(subtree):
+            if child is node:
+                return True
+
+        return False
+
+    @staticmethod
+    def _copy_node(node: ast.AST) -> ast.expr:
+        """Create a shallow copy of an AST node (Attribute or Name)."""
+        if isinstance(node, ast.Name):
+            return cast(ast.expr, ast.Name(id=node.id, ctx=node.ctx))
+        elif isinstance(node, ast.Attribute):
+            value = ModulePropertyTransformer._copy_node(cast(ast.AST, node.value))
+            return cast(ast.expr, ast.Attribute(value=value, attr=node.attr, ctx=node.ctx))
+        return cast(ast.expr, node)

pynecore/transformers/ne_guard.py

@@ -0,0 +1,70 @@
+import ast
+
+
+def _is_skippable_const(node: ast.expr) -> bool:
+    """A ``str``/``bool`` constant operand: Pine only allows comparing it to the
+    same type, so the comparison is homogeneous non-float and a na operand is an
+    ``NA`` object whose ``__ne__`` is already False — no nan guard needed."""
+    return isinstance(node, ast.Constant) and isinstance(node.value, (str, bool))
+
+
+class NeGuardTransformer(ast.NodeTransformer):
+    """
+    Give ``!=`` TradingView's na semantics on native nan operands.
+
+    Pine's float na is a native IEEE-754 nan. Raw IEEE keeps ``==``/``<``/``>``/
+    ``<=``/``>=`` falsy on nan — matching TradingView, where every comparison
+    with na is false — but ``nan != x`` would be True. Compiled scripts keep the
+    readable ``(l) != (r)`` form; this transformer rewrites it at load time to
+
+        l == l and r == r and l != r
+
+    which is False whenever either operand is nan. It is semantically neutral
+    for every other operand type: ``x == x`` is True for str/int/bool/objects,
+    and an ``NA`` object fails its own ``__eq__`` exactly like its ``__ne__``.
+
+    Operands that are not simple names/constants are bound once via a walrus so
+    side effects (function calls) don't run twice. ``str``/``bool`` constant
+    comparisons are homogeneous non-float in Pine and stay untouched.
+    """
+
+    def __init__(self):
+        self._temp_counter = 0
+
+    @staticmethod
+    def _copy_simple(node: ast.expr) -> ast.expr:
+        """Fresh node for re-reading a simple operand (AST nodes must not be shared)."""
+        if isinstance(node, ast.Name):
+            return ast.Name(id=node.id, ctx=ast.Load())
+        assert isinstance(node, ast.Constant)
+        return ast.Constant(value=node.value)
+
+    def _bind_once(self, node: ast.expr) -> tuple[ast.expr, ast.expr, ast.expr]:
+        """Return (first_use, second_use, third_use) for an operand.
+
+        Simple names/constants are re-read; anything else is bound once via a
+        walrus so side effects (function calls) don't run twice."""
+        if isinstance(node, (ast.Name, ast.Constant)):
+            return node, self._copy_simple(node), self._copy_simple(node)
+        self._temp_counter += 1
+        name = f"__ne{self._temp_counter}__"
+        return (ast.NamedExpr(target=ast.Name(id=name, ctx=ast.Store()), value=node),
+                ast.Name(id=name, ctx=ast.Load()),
+                ast.Name(id=name, ctx=ast.Load()))
+
+    def visit_Compare(self, node: ast.Compare) -> ast.expr:
+        self.generic_visit(node)
+
+        if len(node.ops) != 1 or not isinstance(node.ops[0], ast.NotEq):
+            return node
+        left, right = node.left, node.comparators[0]
+        if _is_skippable_const(left) or _is_skippable_const(right):
+            return node
+
+        left1, left2, left3 = self._bind_once(left)
+        right1, right2, right3 = self._bind_once(right)
+        return ast.BoolOp(op=ast.And(), values=[
+            ast.Compare(left=left1, ops=[ast.Eq()], comparators=[left2]),
+            ast.Compare(left=right1, ops=[ast.Eq()], comparators=[right2]),
+            ast.Compare(left=left3, ops=[ast.NotEq()], comparators=[right3]),
+        ])

pynecore/transformers/persistent.py

@@ -0,0 +1,320 @@
+"""
+Transform Persistent type annotations and accesses into state-vector slots.
+
+A persistent variable lives in a compile-time-assigned slot of its scope's
+state vector (a plain list the function receives as its hidden first
+parameter, see :mod:`pynecore.transformers.slot_layout`); every access is a
+literal-index subscript:
+
+- ``p: Persistent[float] = 0.0`` -> slot allocated, declaration removed (the
+  initial value moves into the layout's init template),
+- non-literal initializers keep the lazy pattern: a value slot plus a flag
+  slot and an ``if not __state__[flag]: ...`` guard at the declaration site,
+- reads/writes become ``__state__[N]``,
+- ``+=`` with a non-literal value keeps Kahan summation, emitted as a
+  four-statement sequence (a slot cannot be the target of a walrus, so the
+  legacy single-expression form does not carry over — statement position is
+  guaranteed because ``+=`` is always a statement); variables declared with
+  a ``str``/``bool`` element type skip Kahan (numeric error compensation
+  would crash string concatenation),
+- a walrus write to a persistent (expression position) is emitted through
+  ``__state__.__setitem__`` inside a tuple expression, preserving the value.
+
+Scope rules mirror the legacy transformer: scopes are the middle-dot joined
+function-name path, nested definitions see parent persistents (resolved to
+the parent's state vector through a closure reference on the parent's
+scope-qualified state parameter), and a plain local assignment shadows a
+parent persistent of the same name.
+"""
+from typing import cast
+import ast
+
+from .slot_layout import ModuleLayout, scope_for_function
+
+__all__ = ['PersistentTransformer']
+
+PERSISTENT_TYPES = ('Persistent', 'IBPersistent', 'IBPersistentSeries')
+VARIP_TYPES = ('IBPersistent', 'IBPersistentSeries')
+
+
+class PersistentTransformer(ast.NodeTransformer):
+    """Rewrite Persistent declarations and accesses to state-vector slots."""
+
+    def __init__(self, layout: ModuleLayout):
+        self.layout = layout
+        self.scope_stack: list[str] = []
+        self.current_scope: str = ''
+        # scope -> var name -> (value slot, flag slot or None, kahan slot or None)
+        self.var_slots: dict[str, dict[str, int]] = {}
+        self.kahan_slots: dict[str, dict[str, int]] = {}
+        # scope -> var name -> declared element type name (None when unknown)
+        self.var_types: dict[str, dict[str, str | None]] = {}
+        self.persistent_declarations: dict[str, set[str]] = {}
+        self.local_vars: dict[str, set[str]] = {}
+
+    # --- helpers ---------------------------------------------------------
+
+    def _lookup(self, var_name: str) -> tuple[str, int] | None:
+        """Resolve a name to its declaring scope and value slot.
+
+        A name locally declared in the current scope (but not as Persistent)
+        shadows any parent persistent of the same name.
+
+        :param var_name: Source-level variable name.
+        :return: (declaring scope, slot index) or None.
+        """
+        if (var_name in self.local_vars.get(self.current_scope, ())
+                and var_name not in self.persistent_declarations.get(self.current_scope, ())):
+            return None
+        slots = self.var_slots.get(self.current_scope)
+        if slots is not None and var_name in slots:
+            return self.current_scope, slots[var_name]
+        for i in range(len(self.scope_stack) - 1, 0, -1):
+            scope = '·'.join(self.scope_stack[:i])
+            slots = self.var_slots.get(scope)
+            if slots is not None and var_name in slots:
+                return scope, slots[var_name]
+        return None
+
+    def _state_ref(self, scope: str, slot: int, ctx: ast.expr_context) -> ast.Subscript:
+        """Build a ``<state param>[slot]`` reference for a scope."""
+        return ast.Subscript(
+            value=ast.Name(id=self.layout.state_param(scope), ctx=ast.Load()),
+            slice=ast.Constant(value=slot), ctx=ctx)
+
+    @staticmethod
+    def _is_persistent_type(annotation: ast.expr) -> bool:
+        """Check if the annotation is any form of Persistent type."""
+        if isinstance(annotation, ast.Name):
+            return annotation.id in PERSISTENT_TYPES
+        if isinstance(annotation, ast.Subscript) and isinstance(annotation.value, ast.Name):
+            return annotation.value.id in PERSISTENT_TYPES
+        if isinstance(annotation, ast.Attribute):
+            return annotation.attr in PERSISTENT_TYPES
+        return False
+
+    @staticmethod
+    def _is_varip_type(annotation: ast.expr) -> bool:
+        """Check if the annotation is a varip (IBPersistent) type."""
+        if isinstance(annotation, ast.Name):
+            return annotation.id in VARIP_TYPES
+        if isinstance(annotation, ast.Subscript) and isinstance(annotation.value, ast.Name):
+            return annotation.value.id in VARIP_TYPES
+        if isinstance(annotation, ast.Attribute):
+            return annotation.attr in VARIP_TYPES
+        return False
+
+    @staticmethod
+    def _is_literal_or_na(node: ast.expr) -> bool:
+        """Check if a node is a literal value or ``na``."""
+        if isinstance(node, ast.Constant):
+            return True
+        return isinstance(node, ast.Name) and node.id == 'na'
+
+    @staticmethod
+    def _inner_type_name(annotation: ast.expr) -> str | None:
+        """Declared element type name of a ``Persistent[...]`` annotation."""
+        if isinstance(annotation, ast.Subscript) and isinstance(annotation.slice, ast.Name):
+            return annotation.slice.id
+        return None
+
+    # --- visitors --------------------------------------------------------
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.ImportFrom | None:
+        """Strip the Persistent name from pynecore imports."""
+        if node.module and node.module.startswith('pynecore'):
+            new_names = [name for name in node.names if name.name != 'Persistent']
+            if not new_names:
+                return None
+            node.names = new_names
+        return node
+
+    def visit_Module(self, node: ast.Module) -> ast.Module:
+        self.layout.assign_scope_ids(node)
+        return cast(ast.Module, self.generic_visit(node))
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
+        """Track scopes, qualify the state parameter when nested defs exist."""
+        self.scope_stack.append(self.layout.scope_segment(node))
+        self.current_scope = '·'.join(self.scope_stack)
+
+        scope_for_function(self.layout, self.current_scope, node)
+
+        self.local_vars.setdefault(self.current_scope, set())
+        self.persistent_declarations.setdefault(self.current_scope, set())
+        for arg in node.args.args:
+            self.local_vars[self.current_scope].add(arg.arg)
+
+        node = cast(ast.FunctionDef, self.generic_visit(node))
+
+        # Persistent names are not closure variables anymore — drop them from
+        # nonlocal statements (shadowed locals keep theirs).
+        ancestors = set()
+        for i in range(len(self.scope_stack) - 1, 0, -1):
+            scope = '·'.join(self.scope_stack[:i])
+            ancestors.update(self.persistent_declarations.get(scope, ()))
+        new_body: list[ast.stmt] = []
+        for stmt in node.body:
+            if isinstance(stmt, ast.Nonlocal):
+                stmt.names = [name for name in stmt.names if name not in ancestors]
+                if not stmt.names:
+                    continue
+            new_body.append(stmt)
+        node.body = new_body
+
+        self.scope_stack.pop()
+        self.current_scope = '·'.join(self.scope_stack)
+        return node
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AST | None:
+        """Convert Persistent declarations into slot allocations."""
+        if not (isinstance(node.target, ast.Name) and self._is_persistent_type(node.annotation)):
+            if isinstance(node.target, ast.Name) and self.current_scope:
+                # An annotated assignment declares a local — it shadows a
+                # same-named parent persistent, like a plain assignment does.
+                self.local_vars.setdefault(self.current_scope, set()).add(node.target.id)
+            if node.value:
+                node.value = cast(ast.expr, self.visit(node.value))
+            return node
+
+        if not self.current_scope:
+            raise SyntaxError("Persistent variables must be declared inside a function")
+
+        var_name = node.target.id
+        self.persistent_declarations[self.current_scope].add(var_name)
+        self.local_vars[self.current_scope].add(var_name)
+        self.var_types.setdefault(self.current_scope, {})[var_name] = \
+            self._inner_type_name(node.annotation)
+        varip = self._is_varip_type(node.annotation)
+        scope_layout = self.layout.scope(self.current_scope)
+
+        if node.value is not None and not self._is_literal_or_na(node.value):
+            # Lazy pattern: value slot + flag slot, initializer runs on first call
+            slot = scope_layout.add_var(var_name, ast.Constant(value=None), varip=varip)
+            self.var_slots.setdefault(self.current_scope, {})[var_name] = slot
+            flag = scope_layout.add_flag(var_name)
+            value = cast(ast.expr, self.visit(node.value))
+            return ast.If(
+                test=ast.UnaryOp(op=ast.Not(),
+                                 operand=self._state_ref(self.current_scope, flag, ast.Load())),
+                body=[
+                    ast.Assign(targets=[self._state_ref(self.current_scope, slot, ast.Store())],
+                               value=value),
+                    ast.Assign(targets=[self._state_ref(self.current_scope, flag, ast.Store())],
+                               value=ast.Constant(value=True)),
+                ],
+                orelse=[])
+
+        # A bare-``na`` default must be the na VALUE, not the ``na`` name — at module
+        # level ``na`` resolves to the ``is_na`` function object. Emit ``lib.na()``.
+        na_init = ast.Call(
+            func=ast.Attribute(value=ast.Name(id='lib', ctx=ast.Load()), attr='na', ctx=ast.Load()),
+            args=[], keywords=[])
+        init = node.value if node.value is not None else na_init
+        slot = scope_layout.add_var(var_name, init, varip=varip)
+        self.var_slots.setdefault(self.current_scope, {})[var_name] = slot
+        return None
+
+    def visit_Assign(self, node: ast.Assign) -> ast.Assign:
+        """Convert assignments to persistent variables into slot writes."""
+        if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
+            var_name = cast(ast.Name, node.targets[0]).id
+            # The first plain assignment to a not-yet-known name declares a
+            # local, shadowing a same-named parent persistent.
+            if var_name not in self.local_vars.get(self.current_scope, ()):
+                self.local_vars.setdefault(self.current_scope, set()).add(var_name)
+            found = self._lookup(var_name)
+            if found:
+                scope, slot = found
+                return ast.Assign(
+                    targets=[self._state_ref(scope, slot, ast.Store())],
+                    value=cast(ast.expr, self.visit(node.value)))
+        node.targets = [cast(ast.expr, self.visit(t)) for t in node.targets]
+        node.value = cast(ast.expr, self.visit(node.value))
+        return node
+
+    def visit_AugAssign(self, node: ast.AugAssign) -> ast.AST | list[ast.stmt]:
+        """Slot-target augmented assignment; Kahan summation for non-literal ``+=``."""
+        if isinstance(node.target, ast.Name):
+            found = self._lookup(node.target.id)
+            if found:
+                scope, slot = found
+                if (isinstance(node.op, ast.Add) and not self._is_literal_or_na(node.value)
+                        # Kahan is numeric error compensation — a declared
+                        # str/bool element type concatenates/accumulates plain
+                        and self.var_types.get(scope, {}).get(node.target.id)
+                        not in ('str', 'bool')):
+                    return self._emit_kahan(scope, slot, node.target.id,
+                                            cast(ast.expr, self.visit(node.value)))
+                node.target = self._state_ref(scope, slot, ast.Store())
+                node.value = cast(ast.expr, self.visit(node.value))
+                return node
+        return cast(ast.AST, self.generic_visit(node))
+
+    def _emit_kahan(self, scope: str, slot: int, var_name: str,
+                    value: ast.expr) -> list[ast.stmt]:
+        """Emit the Kahan-summation statement sequence for ``var += value``."""
+        scope_kahans = self.kahan_slots.setdefault(scope, {})
+        comp = scope_kahans.get(var_name)
+        if comp is None:
+            scope_layout = self.layout.scope(scope)
+            comp = scope_kahans[var_name] = scope_layout.add_kahan(
+                var_name, varip=scope_layout.slots[slot].varip)
+
+        def var_ref(ctx: ast.expr_context) -> ast.Subscript:
+            return self._state_ref(scope, slot, ctx)
+
+        def comp_ref(ctx: ast.expr_context) -> ast.Subscript:
+            return self._state_ref(scope, comp, ctx)
+
+        corrected = ast.Name(id='__kahan_corrected__', ctx=ast.Load())
+        new_sum = ast.Name(id='__kahan_new_sum__', ctx=ast.Load())
+        return [
+            # __kahan_corrected__ = <value> - <comp>
+            ast.Assign(targets=[ast.Name(id='__kahan_corrected__', ctx=ast.Store())],
+                       value=ast.BinOp(left=value, op=ast.Sub(), right=comp_ref(ast.Load()))),
+            # __kahan_new_sum__ = <var> + __kahan_corrected__
+            ast.Assign(targets=[ast.Name(id='__kahan_new_sum__', ctx=ast.Store())],
+                       value=ast.BinOp(left=var_ref(ast.Load()), op=ast.Add(), right=corrected)),
+            # <comp> = (__kahan_new_sum__ - <var>) - __kahan_corrected__
+            ast.Assign(targets=[comp_ref(ast.Store())],
+                       value=ast.BinOp(
+                           left=ast.BinOp(left=new_sum, op=ast.Sub(),
+                                          right=var_ref(ast.Load())),
+                           op=ast.Sub(), right=corrected)),
+            # <var> = __kahan_new_sum__
+            ast.Assign(targets=[var_ref(ast.Store())], value=new_sum),
+        ]
+
+    def visit_NamedExpr(self, node: ast.NamedExpr) -> ast.AST:
+        """Walrus write to a persistent: route through ``__setitem__`` so the
+        construct stays a valid expression and still yields the value."""
+        if isinstance(node.target, ast.Name):
+            found = self._lookup(node.target.id)
+            if found:
+                scope, slot = found
+                value = cast(ast.expr, self.visit(node.value))
+                param = self.layout.state_param(scope)
+                return ast.Subscript(
+                    value=ast.Tuple(
+                        elts=[
+                            ast.NamedExpr(target=ast.Name(id='__pyne_w__', ctx=ast.Store()),
+                                          value=value),
+                            ast.Call(
+                                func=ast.Attribute(value=ast.Name(id=param, ctx=ast.Load()),
+                                                   attr='__setitem__', ctx=ast.Load()),
+                                args=[ast.Constant(value=slot),
+                                      ast.Name(id='__pyne_w__', ctx=ast.Load())],
+                                keywords=[]),
+                        ],
+                        ctx=ast.Load()),
+                    slice=ast.Constant(value=0), ctx=ast.Load())
+        return cast(ast.AST, self.generic_visit(node))
+
+    def visit_Name(self, node: ast.Name) -> ast.AST:
+        """Convert persistent references using scope-aware lookup."""
+        found = self._lookup(node.id)
+        if found:
+            scope, slot = found
+            return self._state_ref(scope, slot, node.ctx)
+        return node