sonolus.py 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

sonolus/backend/visitor.py

@@ -1,23 +1,25 @@
 # ruff: noqa: N802
 import ast
+import builtins
 import functools
 import inspect
-from collections.abc import Callable, Mapping
-from contextlib import contextmanager
-from types import FunctionType, MethodType
-from typing import Any, Never
+from collections.abc import Callable, Sequence
+from types import FunctionType, MethodType, MethodWrapperType
+from typing import Any, Never, Self
 
 from sonolus.backend.excepthook import install_excepthook
 from sonolus.backend.utils import get_function, scan_writes
 from sonolus.script.debug import assert_true
-from sonolus.script.internal.builtin_impls import BUILTIN_IMPLS
+from sonolus.script.internal.builtin_impls import BUILTIN_IMPLS, _bool, _float, _int, _len
+from sonolus.script.internal.constant import ConstantValue
 from sonolus.script.internal.context import Context, EmptyBinding, Scope, ValueBinding, ctx, set_ctx
 from sonolus.script.internal.descriptor import SonolusDescriptor
 from sonolus.script.internal.error import CompilationError
-from sonolus.script.internal.impl import try_validate_value, validate_value
+from sonolus.script.internal.impl import validate_value
+from sonolus.script.internal.tuple_impl import TupleImpl
 from sonolus.script.internal.value import Value
 from sonolus.script.iterator import SonolusIterator
-from sonolus.script.num import Num, is_num
+from sonolus.script.num import Num, _is_num
 
 _compiler_internal_ = True
 
@@ -31,7 +33,7 @@ def compile_and_call[**P, R](fn: Callable[P, R], /, *args: P.args, **kwargs: P.k
 def generate_fn_impl(fn: Callable):
     install_excepthook()
     match fn:
-        case Value() as value if value._is_py_():
+        case ConstantValue() as value if value._is_py_():
             return generate_fn_impl(value._as_py_())
         case MethodType() as method:
             return functools.partial(generate_fn_impl(method.__func__), method.__self__)
@@ -42,20 +44,21 @@ def generate_fn_impl(fn: Callable):
         case _:
             if callable(fn) and isinstance(fn, Value):
                 return generate_fn_impl(fn.__call__)
-            elif fn is type:
-                return fn
             elif callable(fn):
                 raise TypeError(f"Unsupported callable {fn!r}")
             else:
-                raise TypeError(f"Not callable {fn!r}")
+                raise TypeError(f"'{type(fn).__name__}' object is not callable")
 
 
 def eval_fn(fn: Callable, /, *args, **kwargs):
     source_file, node = get_function(fn)
     bound_args = inspect.signature(fn).bind(*args, **kwargs)
     bound_args.apply_defaults()
-    closurevars = inspect.getclosurevars(fn)
-    global_vars = {**closurevars.nonlocals, **closurevars.globals, **closurevars.builtins}
+    global_vars = {
+        **builtins.__dict__,
+        **fn.__globals__,
+        **inspect.getclosurevars(fn).nonlocals,
+    }
     return Visitor(source_file, bound_args, global_vars).run(node)
 
 
@@ -123,16 +126,45 @@ comp_ops = {
 }
 
 rcomp_ops = {
-    ast.Eq: "__req__",
-    ast.NotEq: "__rne__",
+    ast.Eq: "__eq__",
+    ast.NotEq: "__ne__",
     ast.Lt: "__gt__",
     ast.LtE: "__ge__",
     ast.Gt: "__lt__",
     ast.GtE: "__le__",
-    ast.In: "__contains__",
+    ast.In: "__contains__",  # Only supported on the right side
     ast.NotIn: "__contains__",
 }
 
+op_to_symbol = {
+    ast.Add: "+",
+    ast.Sub: "-",
+    ast.Mult: "*",
+    ast.Div: "/",
+    ast.FloorDiv: "//",
+    ast.Mod: "%",
+    ast.Pow: "**",
+    ast.Eq: "==",
+    ast.NotEq: "!=",
+    ast.Lt: "<",
+    ast.LtE: "<=",
+    ast.Gt: ">",
+    ast.GtE: ">=",
+    ast.And: "and",
+    ast.Or: "or",
+    ast.BitAnd: "&",
+    ast.BitOr: "|",
+    ast.BitXor: "^",
+    ast.LShift: "<<",
+    ast.RShift: ">>",
+    ast.USub: "-",
+    ast.UAdd: "+",
+    ast.Invert: "~",
+    ast.Not: "not",
+    ast.In: "in",
+    ast.NotIn: "not in",
+}
+
 
 class Visitor(ast.NodeVisitor):
     source_file: str
@@ -142,22 +174,25 @@ class Visitor(ast.NodeVisitor):
     return_ctxs: list[Context]  # Contexts at return statements, which will branch to the exit
     loop_head_ctxs: list[Context]  # Contexts at loop heads, from outer to inner
     break_ctxs: list[list[Context]]  # Contexts at break statements, from outer to inner
-
-    def __init__(self, source_file: str, bound_args: inspect.BoundArguments, global_vars: dict[str, Any]):
+    active_ctx: Context | None  # The active context for use in nested functions=
+    parent: Self | None  # The parent visitor for use in nested functions
+
+    def __init__(
+        self,
+        source_file: str,
+        bound_args: inspect.BoundArguments,
+        global_vars: dict[str, Any],
+        parent: Self | None = None,
+    ):
         self.source_file = source_file
-        self.globals = {}
-        for k, v in global_vars.items():
-            # Unfortunately, inspect.closurevars also includes attributes
-            if v is ctx:
-                raise ValueError("Unexpected use of ctx in non-meta function")
-            value = try_validate_value(BUILTIN_IMPLS.get(id(v), v))
-            if value is not None:
-                self.globals[k] = value
+        self.globals = global_vars
         self.bound_args = bound_args
         self.used_names = {}
         self.return_ctxs = []
         self.loop_head_ctxs = []
         self.break_ctxs = []
+        self.active_ctx = None
+        self.parent = parent
 
     def run(self, node):
         before_ctx = ctx()
@@ -169,17 +204,49 @@ class Visitor(ast.NodeVisitor):
                 ctx().scope.set_value("$return", validate_value(None))
                 for stmt in body:
                     self.visit(stmt)
+            case ast.Lambda(body=body):
+                result = self.visit(body)
+                ctx().scope.set_value("$return", result)
             case _:
                 raise NotImplementedError("Unsupported syntax")
         after_ctx = Context.meet([*self.return_ctxs, ctx()])
+        self.active_ctx = after_ctx
         result_binding = after_ctx.scope.get_binding("$return")
         if not isinstance(result_binding, ValueBinding):
             raise ValueError("Function has conflicting return values")
         set_ctx(after_ctx.branch_with_scope(None, before_ctx.scope.copy()))
         return result_binding.value
 
+    def visit(self, node):
+        """Visit a node."""
+        # We want this here so this is filtered out of tracebacks
+        method = "visit_" + node.__class__.__name__
+        visitor = getattr(self, method, self.generic_visit)
+        with self.reporting_errors_at_node(node):
+            return visitor(node)
+
     def visit_FunctionDef(self, node):
-        raise NotImplementedError("Nested functions are not supported")
+        name = node.name
+        signature = self.arguments_to_signature(node.args)
+
+        def fn(*args, **kwargs):
+            bound = signature.bind(*args, **kwargs)
+            bound.apply_defaults()
+            return Visitor(
+                self.source_file,
+                bound,
+                self.globals,
+                self,
+            ).run(node)
+
+        fn._meta_fn_ = True
+        fn.__name__ = name
+        fn.__qualname__ = name
+
+        for decorator in reversed(node.decorator_list):
+            fn = self.handle_call(decorator, self.visit(decorator), fn)
+
+        ctx().scope.set_value(name, validate_value(fn))
 
     def visit_AsyncFunctionDef(self, node):
         raise NotImplementedError("Async functions are not supported")
@@ -194,7 +261,16 @@ class Visitor(ast.NodeVisitor):
         set_ctx(ctx().into_dead())
 
     def visit_Delete(self, node):
-        raise NotImplementedError("Delete statements are not supported")
+        for target in node.targets:
+            match target:
+                case ast.Name():
+                    raise NotImplementedError("Deleting variables is not supported")
+                case ast.Subscript(value=value, slice=slice):
+                    self.handle_delitem(target, self.visit(value), self.visit(slice))
+                case ast.Attribute():
+                    raise NotImplementedError("Deleting attributes is not supported")
+                case _:
+                    raise NotImplementedError("Unsupported delete target")
 
     def visit_Assign(self, node):
         value = self.visit(node.value)
@@ -215,27 +291,40 @@ class Visitor(ast.NodeVisitor):
             if not self.is_not_implemented(result):
                 if result is not lhs_value:
                     raise ValueError("Inplace operation must return the same object")
-                # Skip the actual assignment because the inplace operation already did the job, as an optimization
-                # There could be side effects of assignment, but that's atypical
+                self.handle_assign(node.target, result)
                 return
         if hasattr(lhs_value, regular_fn_name):
             result = self.handle_call(node, getattr(lhs_value, regular_fn_name), rhs_value)
             if not self.is_not_implemented(result):
                 self.handle_assign(node.target, result)
                 return
-        if hasattr(rhs_value, right_fn_name):
+        if hasattr(rhs_value, right_fn_name) and type(lhs_value) is not type(rhs_value):
             result = self.handle_call(node, getattr(rhs_value, right_fn_name), lhs_value)
             if not self.is_not_implemented(result):
                 self.handle_assign(node.target, result)
                 return
-        raise NotImplementedError("Unsupported augmented assignment")
+        raise TypeError(
+            f"unsupported operand type(s) for {op_to_symbol[type(node.op)]}=: "
+            f"'{type(lhs_value).__name__}' and '{type(rhs_value).__name__}'"
+        )
 
     def visit_AnnAssign(self, node):
         value = self.visit(node.value)
         self.handle_assign(node.target, value)
 
     def visit_For(self, node):
-        iterator = iter(self.visit(node.iter))
+        from sonolus.script.internal.tuple_impl import TupleImpl
+
+        iterable = self.visit(node.iter)
+        if isinstance(iterable, TupleImpl):
+            # Unroll the loop
+            for value in iterable.value:
+                set_ctx(ctx().branch(None))
+                self.handle_assign(node.target, validate_value(value))
+                for stmt in node.body:
+                    self.visit(stmt)
+            return
+        iterator = self.handle_call(node, iterable.__iter__)
         if not isinstance(iterator, SonolusIterator):
             raise ValueError("Unsupported iterator")
         writes = scan_writes(node)
@@ -244,6 +333,11 @@ class Visitor(ast.NodeVisitor):
         self.break_ctxs.append([])
         set_ctx(header_ctx)
         has_next = self.ensure_boolean_num(self.handle_call(node, iterator.has_next))
+        if has_next._is_py_() and not has_next._as_py_():
+            # The loop will never run, continue after evaluating the condition
+            for stmt in node.orelse:
+                self.visit(stmt)
+            return
         ctx().test = has_next.ir()
         body_ctx = ctx().branch(None)
         else_ctx = ctx().branch(0)
@@ -271,6 +365,11 @@ class Visitor(ast.NodeVisitor):
         self.break_ctxs.append([])
         set_ctx(header_ctx)
         test = self.ensure_boolean_num(self.visit(node.test))
+        if test._is_py_() and not test._as_py_():
+            # The loop will never run, continue after evaluating the condition
+            for stmt in node.orelse:
+                self.visit(stmt)
+            return
         ctx().test = test.ir()
         body_ctx = ctx().branch(None)
         else_ctx = ctx().branch(0)
@@ -343,24 +442,38 @@ class Visitor(ast.NodeVisitor):
                         self.visit(stmt)
                     end_ctxs.append(ctx())
                 else:
-                    end_ctxs.append(ctx())
+                    # Merge failing before the guard and failing now at the guard (which we know is guaranteed to fail)
+                    false_ctx = Context.meet([ctx(), false_ctx])
             else:
                 ctx().test = guard.ir()
-                true_ctx = ctx().branch(None)
-                false_ctx = ctx().branch(0)
-                set_ctx(true_ctx)
+                guard_true_ctx = ctx().branch(None)
+                guard_false_ctx = ctx().branch(0)
+                set_ctx(guard_true_ctx)
                 for stmt in case.body:
                     self.visit(stmt)
                 end_ctxs.append(ctx())
+                false_ctx = Context.meet([false_ctx, guard_false_ctx])
             set_ctx(false_ctx)
+        end_ctxs.append(ctx())
         if end_ctxs:
             set_ctx(Context.meet(end_ctxs))
 
     def handle_match_pattern(self, subject: Value, pattern: ast.pattern) -> tuple[Context, Context]:
+        from sonolus.script.internal.generic import validate_type_spec
+        from sonolus.script.internal.tuple_impl import TupleImpl
+
+        if not ctx().live:
+            return ctx().into_dead(), ctx()
+
         match pattern:
             case ast.MatchValue(value=value):
                 value = self.visit(value)
-                test = self.ensure_boolean_num(subject == value)
+                test = self.ensure_boolean_num(validate_value(subject == value))
+                if test._is_py_():
+                    if test._as_py_():
+                        return ctx(), ctx().into_dead()
+                    else:
+                        return ctx().into_dead(), ctx()
                 ctx_init = ctx()
                 ctx_init.test = test.ir()
                 true_ctx = ctx_init.branch(None)
@@ -369,11 +482,11 @@ class Visitor(ast.NodeVisitor):
             case ast.MatchSingleton(value=value):
                 match value:
                     case True:
-                        test = self.ensure_boolean_num(subject)
+                        raise NotImplementedError("Matching against True is not supported, use 1 instead")
                     case False:
-                        test = self.ensure_boolean_num(subject).not_()
+                        raise NotImplementedError("Matching against False is not supported, use 0 instead")
                     case None:
-                        test = Num._accept_(subject._is_py_() and subject._as_py_() is None)
+                        test = validate_value(subject._is_py_() and subject._as_py_() is None)
                     case _:
                         raise NotImplementedError("Unsupported match singleton")
                 ctx_init = ctx()
@@ -381,19 +494,39 @@ class Visitor(ast.NodeVisitor):
                 true_ctx = ctx_init.branch(None)
                 false_ctx = ctx_init.branch(0)
                 return true_ctx, false_ctx
-            case ast.MatchSequence():
-                raise NotImplementedError("Match sequences are not supported")
+            case ast.MatchSequence(patterns=patterns):
+                target_len = len(patterns)
+                if not (isinstance(subject, Sequence | TupleImpl)):
+                    return ctx().into_dead(), ctx()
+                length_test = self.ensure_boolean_num(validate_value(_len(subject) == target_len))
+                ctx_init = ctx()
+                if not length_test._is_py_():
+                    ctx_init.test = length_test.ir()
+                    true_ctx = ctx_init.branch(None)
+                    false_ctxs = [ctx_init.branch(0)]
+                elif length_test._as_py_():
+                    true_ctx = ctx_init
+                    false_ctxs = []
+                else:
+                    return ctx().into_dead(), ctx()
+                set_ctx(true_ctx)
+                for i, subpattern in enumerate(patterns):
+                    if not ctx().live:
+                        break
+                    value = self.handle_getitem(subpattern, subject, validate_value(i))
+                    true_ctx, false_ctx = self.handle_match_pattern(value, subpattern)
+                    false_ctxs.append(false_ctx)
+                    set_ctx(true_ctx)
+                return true_ctx, Context.meet(false_ctxs)
             case ast.MatchMapping():
                 raise NotImplementedError("Match mappings are not supported")
             case ast.MatchClass(cls=cls, patterns=patterns, kwd_attrs=kwd_attrs, kwd_patterns=kwd_patterns):
-                from sonolus.script.comptime import Comptime
-                from sonolus.script.internal.generic import validate_type_spec
-
-                cls = validate_type_spec(self.visit(cls))
+                cls = self.visit(cls)
+                if cls._is_py_() and cls._as_py_() in {_int, _float, _bool}:
+                    raise TypeError("Instance check against int, float, or bool is not supported, use Num instead")
+                cls = validate_type_spec(cls)
                 if not isinstance(cls, type):
                     raise TypeError("Class is not a type")
-                if issubclass(cls, Comptime):
-                    raise TypeError("Comptime is not supported in match patterns")
                 if not isinstance(subject, cls):
                     return ctx().into_dead(), ctx()
                 if patterns:
@@ -431,13 +564,14 @@ class Visitor(ast.NodeVisitor):
                     return ctx(), ctx().into_dead()
             case ast.MatchOr():
                 true_ctxs = []
-                false_ctx = ctx()
                 assert pattern.patterns
                 for subpattern in pattern.patterns:
+                    if not ctx().live:
+                        break
                     true_ctx, false_ctx = self.handle_match_pattern(subject, subpattern)
                     true_ctxs.append(true_ctx)
                     set_ctx(false_ctx)
-                return Context.meet(true_ctxs), false_ctx
+                return Context.meet(true_ctxs), ctx()
 
     def visit_Raise(self, node):
         raise NotImplementedError("Raise statements are not supported")
@@ -486,7 +620,7 @@ class Visitor(ast.NodeVisitor):
             case ast.Or():
                 handler = self.handle_or
             case _:
-                raise NotImplementedError(f"Unsupported bool operator {node.op}")
+                raise NotImplementedError(f"Unsupported bool operator {op_to_symbol[type(node.op)]}")
 
         if not node.values:
             raise ValueError("Bool operator requires at least one operand")
@@ -504,15 +638,23 @@ class Visitor(ast.NodeVisitor):
         lhs = self.visit(node.left)
         rhs = self.visit(node.right)
         op = bin_ops[type(node.op)]
+        if lhs._is_py_() and rhs._is_py_():
+            lhs_py = lhs._as_py_()
+            rhs_py = rhs._as_py_()
+            if isinstance(lhs_py, type) and isinstance(rhs_py, type):
+                return validate_value(getattr(lhs_py, op)(rhs_py))
         if hasattr(lhs, op):
             result = self.handle_call(node, getattr(lhs, op), rhs)
             if not self.is_not_implemented(result):
                 return result
-        if hasattr(rhs, rbin_ops[type(node.op)]):
+        if hasattr(rhs, rbin_ops[type(node.op)]) and type(lhs) is not type(rhs):
             result = self.handle_call(node, getattr(rhs, rbin_ops[type(node.op)]), lhs)
             if not self.is_not_implemented(result):
                 return result
-        raise NotImplementedError(f"Unsupported operand types for binary operator {node.op}")
+        raise TypeError(
+            f"unsupported operand type(s) for {op_to_symbol[type(node.op)]}: "
+            f"'{type(lhs).__name__}' and '{type(rhs).__name__}'"
+        )
 
     def visit_UnaryOp(self, node):
         operand = self.visit(node.operand)
@@ -521,10 +663,25 @@ class Visitor(ast.NodeVisitor):
         op = unary_ops[type(node.op)]
         if hasattr(operand, op):
             return self.handle_call(node, getattr(operand, op))
-        raise NotImplementedError(f"Unsupported operand type for unary operator {node.op}")
+        raise TypeError(f"bad operand type for unary {op_to_symbol[type(node.op)]}: '{type(operand).__name__}'")
 
     def visit_Lambda(self, node):
-        raise NotImplementedError("Lambda functions are not supported")
+        signature = self.arguments_to_signature(node.args)
+
+        def fn(*args, **kwargs):
+            bound = signature.bind(*args, **kwargs)
+            bound.apply_defaults()
+            return Visitor(
+                self.source_file,
+                bound,
+                self.globals,
+                self,
+            ).run(node)
+
+        fn._meta_fn_ = True
+        fn.__name__ = "<lambda>"
+
+        return validate_value(fn)
 
     def visit_IfExp(self, node):
         test = self.ensure_boolean_num(self.visit(node.test))
@@ -579,6 +736,9 @@ class Visitor(ast.NodeVisitor):
     def visit_YieldFrom(self, node):
         raise NotImplementedError("Yield from expressions are not supported")
 
+    def _has_real_method(self, obj: Value, method_name: str) -> bool:
+        return hasattr(obj, method_name) and not isinstance(getattr(obj, method_name), MethodWrapperType)
+
     def visit_Compare(self, node):
         result_name = self.new_name("compare")
         ctx().scope.set_value(result_name, Num._accept_(0))
@@ -595,22 +755,37 @@ class Visitor(ast.NodeVisitor):
                     result = Num._accept_(l_val._is_py_() and l_val._as_py_() is None)
                 else:
                     result = Num._accept_(not (l_val._is_py_() and l_val._as_py_() is None))
-            elif type(op) in comp_ops and hasattr(l_val, comp_ops[type(op)]):
+            elif type(op) in comp_ops and self._has_real_method(l_val, comp_ops[type(op)]):
                 result = self.handle_call(node, getattr(l_val, comp_ops[type(op)]), r_val)
             if (
                 (result is None or self.is_not_implemented(result))
                 and type(op) in rcomp_ops
-                and hasattr(r_val, rcomp_ops[type(op)])
+                and self._has_real_method(r_val, rcomp_ops[type(op)])
             ):
                 result = self.handle_call(node, getattr(r_val, rcomp_ops[type(op)]), l_val)
             if result is None or self.is_not_implemented(result):
-                raise NotImplementedError(f"Unsupported comparison operator {op}")
+                if type(op) is ast.Eq:
+                    result = Num._accept_(l_val is r_val)
+                elif type(op) is ast.NotEq:
+                    result = Num._accept_(l_val is not r_val)
+                else:
+                    raise TypeError(
+                        f"'{op_to_symbol[type(op)]}' not supported between instances of '{type(l_val).__name__}' and "
+                        f"'{type(r_val).__name__}'"
+                    )
             result = self.ensure_boolean_num(result)
             if inverted:
                 result = result.not_()
             curr_ctx = ctx()
             if i == len(node.ops) - 1:
                 curr_ctx.scope.set_value(result_name, result)
+            elif result._is_py_():
+                if result._as_py_():
+                    l_val = r_val
+                else:
+                    false_ctxs.append(curr_ctx)
+                    set_ctx(curr_ctx.into_dead())
+                    break
             else:
                 curr_ctx.test = result.ir()
                 true_ctx = curr_ctx.branch(None)
@@ -623,9 +798,9 @@ class Visitor(ast.NodeVisitor):
         return ctx().scope.get_value(result_name)
 
     def visit_Call(self, node):
+        from sonolus.script.internal.dict_impl import DictImpl
+
         fn = self.visit(node.func)
-        if fn is Num:
-            raise ValueError("Calling int/bool/float is not supported")
         args = []
         kwargs = {}
         for arg in node.args:
@@ -638,8 +813,10 @@ class Visitor(ast.NodeVisitor):
                 kwargs[keyword.arg] = self.visit(keyword.value)
             else:
                 value = self.visit(keyword.value)
-                if value._is_py_() and isinstance(value._as_py_(), Mapping):
-                    kwargs.update(value._as_py_())
+                if isinstance(value, DictImpl):
+                    if not all(isinstance(k, str) for k in value.value):
+                        raise ValueError("Keyword arguments must be strings")
+                    kwargs.update(value.value)
                 else:
                     raise ValueError("Starred keyword arguments (**kwargs) must be dictionaries")
         return self.handle_call(node, fn, *args, **kwargs)
@@ -665,11 +842,18 @@ class Visitor(ast.NodeVisitor):
         raise NotImplementedError("Starred expressions are not supported")
 
     def visit_Name(self, node):
-        if isinstance(ctx().scope.get_binding(node.id), EmptyBinding) and node.id in self.globals:
-            # globals can have false positives due to limitations of inspect.closurevars
-            # so we need to check that it's not defined as a local variable
-            return self.globals[node.id]
-        return ctx().scope.get_value(node.id)
+        self.active_ctx = ctx()
+        v = self
+        while v:
+            if not isinstance(v.active_ctx.scope.get_binding(node.id), EmptyBinding):
+                return v.active_ctx.scope.get_value(node.id)
+            v = v.parent
+        if node.id in self.globals:
+            value = self.globals[node.id]
+            if value is ctx:
+                raise ValueError("Unexpected use of ctx in non meta-function")
+            return validate_value(BUILTIN_IMPLS.get(id(value), value))
+        raise NameError(f"Name {node.id} is not defined")
 
     def visit_List(self, node):
         raise NotImplementedError("List literals are not supported")
@@ -712,8 +896,12 @@ class Visitor(ast.NodeVisitor):
         ctx_init = ctx()
         l_val = self.ensure_boolean_num(l_val)
 
-        if l_val._is_py_() and not l_val._as_py_():
-            return l_val
+        if l_val._is_py_():
+            if l_val._as_py_():
+                # The rhs is definitely evaluated, so we can return it directly
+                return self.ensure_boolean_num(self.visit(r_expr))
+            else:
+                return l_val
 
         ctx_init.test = l_val.ir()
         res_name = self.new_name("and")
@@ -736,8 +924,12 @@ class Visitor(ast.NodeVisitor):
         ctx_init = ctx()
         l_val = self.ensure_boolean_num(l_val)
 
-        if l_val._is_py_() and l_val._as_py_():
-            return l_val
+        if l_val._is_py_():
+            if l_val._as_py_():
+                return l_val
+            else:
+                # The rhs is definitely evaluated, so we can return it directly
+                return self.ensure_boolean_num(self.visit(r_expr))
 
         ctx_init.test = l_val.ir()
         res_name = self.new_name("or")
@@ -764,7 +956,8 @@ class Visitor(ast.NodeVisitor):
 
     def handle_getattr(self, node: ast.stmt | ast.expr, target: Value, key: str) -> Value:
         with self.reporting_errors_at_node(node):
-            if target._is_py_():
+            if isinstance(target, ConstantValue):
+                # Unwrap so we can access fields
                 target = target._as_py_()
             descriptor = type(target).__dict__.get(key)
             match descriptor:
@@ -796,6 +989,7 @@ class Visitor(ast.NodeVisitor):
         self, node: ast.stmt | ast.expr, fn: Callable[P, R], /, *args: P.args, **kwargs: P.kwargs
     ) -> R:
         """Handles a call to the given callable."""
+        self.active_ctx = ctx()
         if (
             isinstance(fn, Value)
             and fn._is_py_()
@@ -808,13 +1002,10 @@ class Visitor(ast.NodeVisitor):
 
     def handle_getitem(self, node: ast.stmt | ast.expr, target: Value, key: Value) -> Value:
         with self.reporting_errors_at_node(node):
-            if target._is_py_():
-                target = target._as_py_()
-                if key._is_py_():
-                    return validate_value(target[key._as_py_()])
-                if isinstance(target, Value) and hasattr(target, "__getitem__"):
-                    return self.handle_call(node, target.__getitem__, key)
-                raise TypeError(f"Cannot get items on {type(target).__name__}")
+            if target._is_py_() and isinstance(target._as_py_(), type):
+                if not key._is_py_():
+                    raise ValueError("Type parameters must be compile-time constants")
+                return validate_value(target._as_py_()[key._as_py_()])
             else:
                 if isinstance(target, Value) and hasattr(target, "__getitem__"):
                     return self.handle_call(node, target.__getitem__, key)
@@ -822,21 +1013,19 @@ class Visitor(ast.NodeVisitor):
 
     def handle_setitem(self, node: ast.stmt | ast.expr, target: Value, key: Value, value: Value):
         with self.reporting_errors_at_node(node):
-            if target._is_py_():
-                target = target._as_py_()
-                if key._is_py_():
-                    target[key._as_py_()] = value._as_py_()
-                if isinstance(target, Value) and hasattr(target, "__setitem__"):
-                    return self.handle_call(node, target.__setitem__, key, value)
-                raise TypeError(f"Cannot set items on {type(target).__name__}")
-            else:
-                if isinstance(target, Value) and hasattr(target, "__setitem__"):
-                    return self.handle_call(node, target.__setitem__, key, value)
-                raise TypeError(f"Cannot set items on {type(target).__name__}")
+            if isinstance(target, Value) and hasattr(target, "__setitem__"):
+                return self.handle_call(node, target.__setitem__, key, value)
+            raise TypeError(f"Cannot set items on {type(target).__name__}")
+
+    def handle_delitem(self, node: ast.stmt | ast.expr, target: Value, key: Value):
+        with self.reporting_errors_at_node(node):
+            if isinstance(target, Value) and hasattr(target, "__delitem__"):
+                return self.handle_call(node, target.__delitem__, key)
+            raise TypeError(f"Cannot delete items on {type(target).__name__}")
 
     def handle_starred(self, value: Value) -> tuple[Value, ...]:
-        if value._is_py_() and isinstance(value._as_py_(), tuple):
-            return value._as_py_()
+        if isinstance(value, TupleImpl):
+            return value.value
         raise ValueError("Unsupported starred expression")
 
     def is_not_implemented(self, value):
@@ -845,10 +1034,66 @@ class Visitor(ast.NodeVisitor):
 
     def ensure_boolean_num(self, value) -> Num:
         # This just checks the type for now, although we could support custom __bool__ implementations in the future
-        if not is_num(value):
+        if not _is_num(value):
             raise TypeError(f"Invalid type where a bool (Num) was expected: {type(value).__name__}")
         return value
 
+    def arguments_to_signature(self, arguments: ast.arguments) -> inspect.Signature:
+        parameters: list[inspect.Parameter] = []
+        pos_only_count = len(arguments.posonlyargs)
+        for i, arg in enumerate(arguments.posonlyargs):
+            default_idx = i - pos_only_count + len(arguments.defaults)
+            default = self.visit(arguments.defaults[default_idx]) if default_idx >= 0 else None
+            param = inspect.Parameter(
+                name=arg.arg,
+                kind=inspect.Parameter.POSITIONAL_ONLY,
+                default=default if default_idx >= 0 else inspect.Parameter.empty,
+                annotation=inspect.Parameter.empty,
+            )
+            parameters.append(param)
+
+        pos_kw_count = len(arguments.args)
+        for i, arg in enumerate(arguments.args):
+            default_idx = i - pos_kw_count + len(arguments.defaults)
+            default = self.visit(arguments.defaults[default_idx]) if default_idx >= 0 else None
+            param = inspect.Parameter(
+                name=arg.arg,
+                kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                default=default if default_idx >= 0 else inspect.Parameter.empty,
+                annotation=inspect.Parameter.empty,
+            )
+            parameters.append(param)
+
+        if arguments.vararg:
+            param = inspect.Parameter(
+                name=arguments.vararg.arg,
+                kind=inspect.Parameter.VAR_POSITIONAL,
+                default=inspect.Parameter.empty,
+                annotation=inspect.Parameter.empty,
+            )
+            parameters.append(param)
+
+        for i, arg in enumerate(arguments.kwonlyargs):
+            default = self.visit(arguments.kw_defaults[i]) if arguments.kw_defaults[i] is not None else None
+            param = inspect.Parameter(
+                name=arg.arg,
+                kind=inspect.Parameter.KEYWORD_ONLY,
+                default=default if default is not None else inspect.Parameter.empty,
+                annotation=inspect.Parameter.empty,
+            )
+            parameters.append(param)
+
+        if arguments.kwarg:
+            param = inspect.Parameter(
+                name=arguments.kwarg.arg,
+                kind=inspect.Parameter.VAR_KEYWORD,
+                default=inspect.Parameter.empty,
+                annotation=inspect.Parameter.empty,
+            )
+            parameters.append(param)
+
+        return inspect.Signature(parameters)
+
     def raise_exception_at_node(self, node: ast.stmt | ast.expr, cause: Exception) -> Never:
         """Throws a compilation error at the given node."""
 
@@ -878,15 +1123,29 @@ class Visitor(ast.NodeVisitor):
             {"fn": fn, "args": args, "kwargs": kwargs, "_filter_traceback_": True},
         )
 
-    @contextmanager
     def reporting_errors_at_node(self, node: ast.stmt | ast.expr):
-        try:
-            yield
-        except CompilationError as e:
-            raise e from None
-        except Exception as e:
-            self.raise_exception_at_node(node, e)
+        return ReportingErrorsAtNode(self, node)
 
     def new_name(self, name: str):
         self.used_names[name] = self.used_names.get(name, 0) + 1
         return f"${name}_{self.used_names[name]}"
+
+
+# Not using @contextmanager so it doesn't end up in tracebacks
+class ReportingErrorsAtNode:
+    def __init__(self, compiler, node: ast.stmt | ast.expr):
+        self.compiler = compiler
+        self.node = node
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            return
+
+        if issubclass(exc_type, CompilationError):
+            raise exc_value from exc_value.__cause__
+
+        if exc_value is not None:
+            self.compiler.raise_exception_at_node(self.node, exc_value)