sonolus.py 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

sonolus/script/internal/context.py

@@ -19,6 +19,25 @@ _compiler_internal_ = True
 context_var = ContextVar("context_var", default=None)
 
 
+@dataclass(frozen=True)
+class DebugConfig:
+    unchecked_reads: bool
+    unchecked_writes: bool
+
+
+_full_debug_config = DebugConfig(
+    unchecked_reads=True,
+    unchecked_writes=True,
+)
+
+_disabled_debug_config = DebugConfig(
+    unchecked_reads=False,
+    unchecked_writes=False,
+)
+
+debug_var = ContextVar("debug_var", default=_disabled_debug_config)
+
+
 class GlobalContextState:
     archetypes: dict[type, int]
     rom: ReadOnlyMemory
@@ -92,12 +111,16 @@ class Context:
         return self.callback_state.used_names
 
     def check_readable(self, place: BlockPlace):
+        if debug_config().unchecked_reads:
+            return
         if not self.callback:
             return
         if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).readable:
             raise RuntimeError(f"Block {place.block} is not readable in {self.callback}")
 
     def check_writable(self, place: BlockPlace):
+        if debug_config().unchecked_writes:
+            return
         if not self.callback:
             return
         if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).writable:
@@ -238,6 +261,12 @@ class Context:
             context.outgoing[None] = target
         return target
 
+    def register_archetype(self, type_: type) -> int:
+        with self.global_state.lock:
+            if type_ not in self.global_state.archetypes:
+                self.global_state.archetypes[type_] = len(self.global_state.archetypes)
+            return self.global_state.archetypes[type_]
+
 
 def ctx() -> Context | None:
     return context_var.get()
@@ -262,7 +291,11 @@ class ReadOnlyMemory:
     _lock: Lock
 
     def __init__(self):
-        self.values = []
+        self.values = [
+            float("nan"),
+            float("inf"),
+            float("-inf"),
+        ]
         self.indexes = {}
         self._lock = Lock()
 
@@ -283,6 +316,27 @@ class ReadOnlyMemory:
         else:
             return PlayBlock.EngineRom
 
+    def get_value(self, index: int) -> float:
+        with self._lock:
+            if index < 0 or index >= len(self.values):
+                raise IndexError(f"Index {index} out of bounds for ReadOnlyMemory")
+            return self.values[index]
+
+
+@contextmanager
+def enable_debug(config: DebugConfig | None = None):
+    if config is None:
+        config = _full_debug_config
+    token = debug_var.set(config)
+    try:
+        yield
+    finally:
+        debug_var.reset(token)
+
+
+def debug_config() -> DebugConfig:
+    return debug_var.get()
+
 
 @dataclass
 class ValueBinding:

sonolus/script/internal/range.py

@@ -1,4 +1,6 @@
-from sonolus.script.array_like import ArrayLike
+from sonolus.script.array_like import ArrayLike, get_positive_index
+from sonolus.script.internal.context import ctx
+from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.iterator import SonolusIterator
 from sonolus.script.num import Num
 from sonolus.script.record import Record
@@ -36,7 +38,7 @@ class Range(Record, ArrayLike[Num]):
             return (diff - self.step - 1) // -self.step
 
     def __getitem__(self, index: Num) -> Num:
-        return self.start + index * self.step
+        return self.start + get_positive_index(index, len(self)) * self.step
 
     def __setitem__(self, index: Num, value: Num):
         raise TypeError("Range does not support item assignment")
@@ -79,3 +81,24 @@ class RangeIterator(Record, SonolusIterator):
 
     def advance(self):
         self.value += self.step
+
+
+@meta_fn
+def range_or_tuple(start: Num, stop: Num | None = None, step: Num = 1) -> Range | tuple[Num, ...]:
+    if stop is None:
+        start, stop = 0, start
+    if not ctx():
+        return range(start, stop, step)  # type: ignore
+    start = Num._accept_(start)
+    stop = Num._accept_(stop) if stop is not None else None
+    step = Num._accept_(step)
+    if start._is_py_() and stop._is_py_() and step._is_py_():
+        start_int = start._as_py_()
+        stop_int = stop._as_py_() if stop is not None else None
+        if stop_int is None:
+            start_int, stop_int = 0, start_int
+        step_int = step._as_py_()
+        if start_int % 1 != 0 or stop_int % 1 != 0 or step_int % 1 != 0:
+            raise TypeError("Range arguments must be integers")
+        return validate_value(tuple(range(int(start_int), int(stop_int), int(step_int))))  # type: ignore
+    return Range(start, stop, step)

sonolus/script/internal/simulation_context.py

@@ -0,0 +1,131 @@
+import builtins
+import sys
+from collections.abc import Callable
+from types import MethodType, ModuleType
+from typing import Any, ClassVar, Self
+
+from sonolus.script.internal.impl import validate_value
+
+_MISSING = object()
+
+
+class SimulationContext:
+    """Context manager to simulate additional Sonolus runtime features like level memory."""
+
+    _active_context: ClassVar[Self | None] = None
+    additional_replacements: dict[Any, Any]
+    values: dict[Any, Any]
+    _original_values: dict[str, dict[str, Any]]  # module_name: {var_name: original_value}
+    _original_import: Callable[[str, Any, Any, Any, Any], Any] | None
+
+    def __init__(self, *, additional_replacements: dict[Any, Any] | None = None):
+        if SimulationContext._active_context is not None:
+            raise RuntimeError("SimulationContext is already active")
+        SimulationContext._active_context = self
+        self.additional_replacements = additional_replacements or {}
+        self.values = {}
+        self._original_values = {}
+        self._original_import = None
+
+    def get_or_put_value(self, key: Any, factory: Callable[[], Any]) -> Any:
+        if key not in self.values:
+            self.values[key] = validate_value(factory())
+        return self.values[key]._get_()._as_py_()
+
+    def set_or_put_value(self, key: Any, factory: Callable[[], Any], value: Any):
+        if key not in self.values:
+            self.values[key] = validate_value(factory())
+        existing_value = self.values[key]
+        existing_type = type(existing_value)
+        value = existing_type._accept_(validate_value(value))
+        if existing_type._is_value_type_():
+            existing_value._set_(value)
+        else:
+            existing_value._copy_from_(value)
+
+    def _get_replacement(self, value: Any) -> Any:
+        try:
+            if value in self.additional_replacements:
+                return self.additional_replacements[value]
+        except TypeError:
+            pass
+        if hasattr(value, "_get_sim_replacement_") and isinstance(value._get_sim_replacement_, MethodType):
+            return value._get_sim_replacement_()
+        return _MISSING
+
+    def _substitute_module_variables(self, module) -> None:
+        if not isinstance(module, ModuleType):
+            return
+
+        module_name = module.__name__
+
+        if module_name in self._original_values:
+            return
+
+        original_values = {}
+
+        for var_name, var_value in list(module.__dict__.items()):
+            replacement = self._get_replacement(var_value)
+            if replacement is not _MISSING:
+                original_values[var_name] = var_value
+                setattr(module, var_name, replacement)
+
+        if original_values:
+            self._original_values[module_name] = original_values
+
+    def _update_loaded_modules(self) -> None:
+        loaded_modules = list(sys.modules.values())
+
+        for module in loaded_modules:
+            if module is not None:
+                self._substitute_module_variables(module)
+
+    def _create_import_hook(self):
+        original_import = builtins.__import__
+
+        def hooked_import(name, globals=None, locals=None, fromlist=(), level=0):  # noqa: A002
+            module = original_import(name, globals, locals, fromlist, level)
+            if isinstance(module, ModuleType):
+                self._substitute_module_variables(module)
+
+            if fromlist:
+                for item in fromlist:
+                    if hasattr(module, item):
+                        attr = getattr(module, item)
+                        if isinstance(attr, ModuleType):
+                            self._substitute_module_variables(attr)
+
+            return module
+
+        return hooked_import
+
+    def __enter__(self) -> Self:
+        self._update_loaded_modules()
+
+        self._original_import = builtins.__import__
+        builtins.__import__ = self._create_import_hook()
+
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+        try:
+            if self._original_import is not None:
+                builtins.__import__ = self._original_import
+                self._original_import = None
+
+            for module_name, original_values in self._original_values.items():
+                if module_name in sys.modules:
+                    module = sys.modules[module_name]
+                    if isinstance(module, ModuleType):
+                        for var_name, original_value in original_values.items():
+                            setattr(module, var_name, original_value)
+
+            self._original_values.clear()
+
+        finally:
+            SimulationContext._active_context = None
+
+
+def sim_ctx() -> SimulationContext | None:
+    """Get the current simulation context, or None if not active."""
+    return SimulationContext._active_context

sonolus/script/internal/tuple_impl.py

@@ -21,8 +21,10 @@ class TupleImpl(TransientValue):
             raise TypeError(f"Cannot index tuple with {item}")
         if int(item) != item:
             raise TypeError(f"Cannot index tuple with non-integer {item}")
-        if not (0 <= item < len(self.value)):
+        if not (-len(self.value) <= item < len(self.value)):
             raise IndexError(f"Tuple index out of range: {item}")
+        if item < 0:
+            item += len(self.value)
         return self.value[int(item)]
 
     @meta_fn
@@ -30,7 +32,7 @@ class TupleImpl(TransientValue):
         return len(self.value)
 
     def __eq__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return False
         if len(self) != len(other):
             return False
@@ -40,7 +42,7 @@ class TupleImpl(TransientValue):
         return True
 
     def __ne__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return True
         if len(self) != len(other):
             return True
@@ -50,33 +52,33 @@ class TupleImpl(TransientValue):
         return False
 
     def __lt__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a < b
         return len(self.value) < len(other.value)
 
     def __le__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a < b
         return len(self.value) <= len(other.value)
 
     def __gt__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a > b
         return len(self.value) > len(other.value)
 
     def __ge__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a > b
         return len(self.value) >= len(other.value)
@@ -89,6 +91,11 @@ class TupleImpl(TransientValue):
         other = TupleImpl._accept_(other)
         return TupleImpl._accept_(self.value + other.value)
 
+    @staticmethod
+    @meta_fn
+    def _is_tuple_impl(value: Any) -> bool:
+        return isinstance(value, TupleImpl)
+
     @classmethod
     def _accepts_(cls, value: Any) -> bool:
         return isinstance(value, cls | tuple)

sonolus/script/interval.py

@@ -1,8 +1,10 @@
 from typing import Self
 
 from sonolus.backend.ops import Op
+from sonolus.script.array_like import ArrayLike
 from sonolus.script.debug import static_error
 from sonolus.script.internal.native import native_function
+from sonolus.script.internal.range import range_or_tuple
 from sonolus.script.num import Num
 from sonolus.script.record import Record
 
@@ -296,7 +298,7 @@ def unlerp_clamped(a: float, b: float, x: float, /) -> float:
 
 @native_function(Op.Remap)
 def remap(a: float, b: float, c: float, d: float, x: float, /) -> float:
-    """Remap a value from one interval to another.
+    """Linearly remap a value from one interval to another.
 
     Args:
         a: The start of the input interval.
@@ -313,7 +315,7 @@ def remap(a: float, b: float, c: float, d: float, x: float, /) -> float:
 
 @native_function(Op.RemapClamped)
 def remap_clamped(a: float, b: float, c: float, d: float, x: float, /) -> float:
-    """Remap a value from one interval to another, clamped to the output interval.
+    """Linearly remap a value from one interval to another, clamped to the output interval.
 
     Args:
         a: The start of the input interval.
@@ -341,3 +343,59 @@ def clamp(x: float, a: float, b: float, /) -> float:
         The clamped value.
     """
     return max(a, min(b, x))
+
+
+def interp(
+    xp: ArrayLike[float] | tuple[float, ...],
+    fp: ArrayLike[float] | tuple[float, ...],
+    x: float,
+) -> float:
+    """Linearly interpolate a value within a sequence of points.
+
+    The sequence must have at least 2 elements and be sorted in increasing order of x-coordinates.
+    For values of x outside the range of xp, the slope of the first or last segment is used to extrapolate.
+
+    Args:
+        xp: The x-coordinates of the points in increasing order.
+        fp: The y-coordinates of the points.
+        x: The x-coordinate to interpolate.
+
+    Returns:
+        The interpolated value.
+    """
+    assert len(xp) == len(fp)
+    assert len(xp) >= 2
+    for i in range_or_tuple(1, len(xp) - 1):
+        # At i == 1, x may be less than x[0], but since we're extrapolating, we use the first segment regardless.
+        if x <= xp[i]:
+            return remap(xp[i - 1], xp[i], fp[i - 1], fp[i], x)
+    # x > xp[-2] so we can just use the last segment regardless of whether x is in it or to the right of it.
+    return remap(xp[-2], xp[-1], fp[-2], fp[-1], x)
+
+
+def interp_clamped(
+    xp: ArrayLike[float] | tuple[float, ...],
+    fp: ArrayLike[float] | tuple[float, ...],
+    x: float,
+):
+    """Linearly interpolate a value within a sequence of points.
+
+    The sequence must have at least 2 elements and be sorted in increasing order of x-coordinates.
+    For x-coordinates outside the range of the sequence, the respective endpoint of fp is returned.
+
+    Args:
+        xp: The x-coordinates of the points in increasing order.
+        fp: The y-coordinates of the points.
+        x: The x-coordinate to interpolate.
+
+    Returns:
+        The interpolated value.
+    """
+    assert len(xp) == len(fp)
+    assert len(xp) >= 2
+    if x <= xp[0]:
+        return fp[0]
+    for i in range_or_tuple(1, len(xp)):
+        if x <= xp[i]:
+            return remap(xp[i - 1], xp[i], fp[i - 1], fp[i], x)
+    return fp[-1]

sonolus/script/iterator.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Iterator
 from typing import Any
 
 from sonolus.script.internal.impl import meta_fn
 from sonolus.script.record import Record
 
 
-class SonolusIterator[T](Iterator[T]):
+class SonolusIterator[T]:
     """Base class for Sonolus iterators.
 
     This class is used to define custom iterators that can be used in Sonolus.py.
@@ -30,6 +29,8 @@ class SonolusIterator[T](Iterator[T]):
         ```
     """
 
+    _allow_instance_check_ = True
+
     def next(self) -> T:
         result = self.get()
         self.advance()

sonolus/script/num.py

@@ -1,4 +1,3 @@
-# ruff: noqa: N801
 from __future__ import annotations
 
 import operator
@@ -81,12 +80,22 @@ class _Num(Value, metaclass=_NumMeta):
             return value
         return cls(value)
 
+    def _is_rom_constant(self) -> bool:
+        return (
+            ctx()
+            and isinstance(self.data, BlockPlace)
+            and self.data.block == ctx().blocks.EngineRom
+            and isinstance(self.data.index, int)
+        )
+
     def _is_py_(self) -> bool:
-        return isinstance(self.data, float | int | bool)
+        return isinstance(self.data, float | int | bool) or self._is_rom_constant()
 
     def _as_py_(self) -> Any:
         if not self._is_py_():
             raise ValueError("Not a compile time constant Num")
+        if self._is_rom_constant():
+            return ctx().rom.get_value(self.data.index + self.data.offset)
         if self.data.is_integer():
             return int(self.data)
         return self.data

sonolus/script/options.py

@@ -5,11 +5,13 @@ from typing import Annotated, Any, NewType, dataclass_transform, get_origin
 from sonolus.backend.mode import Mode
 from sonolus.backend.place import BlockPlace
 from sonolus.script.debug import assert_unreachable
-from sonolus.script.internal.context import ctx
+from sonolus.script.internal.context import ctx, debug_config
 from sonolus.script.internal.descriptor import SonolusDescriptor
 from sonolus.script.internal.generic import validate_concrete_type
 from sonolus.script.internal.introspection import get_field_specifiers
+from sonolus.script.internal.simulation_context import sim_ctx
 from sonolus.script.num import Num
+from sonolus.script.values import copy
 
 
 @dataclass
@@ -186,6 +188,8 @@ class _OptionField(SonolusDescriptor):
         self.index = index
 
     def __get__(self, instance, owner):
+        if sim_ctx():
+            return sim_ctx().get_or_put_value((instance, self), lambda: copy(self.info.default))
         if ctx():
             match ctx().global_state.mode:
                 case Mode.PLAY:
@@ -202,8 +206,27 @@ class _OptionField(SonolusDescriptor):
                 return Num._from_place_(BlockPlace(block, self.index))
             else:
                 return Num._accept_(self.info.default)
+        raise RuntimeError("Options can only be accessed in a context")
 
     def __set__(self, instance, value):
+        if sim_ctx():
+            return sim_ctx().set_or_put_value((instance, self), lambda: copy(self.info.default), value)
+        if ctx() and debug_config().unchecked_writes:
+            match ctx().global_state.mode:
+                case Mode.PLAY:
+                    block = ctx().blocks.LevelOption
+                case Mode.WATCH:
+                    block = ctx().blocks.LevelOption
+                case Mode.PREVIEW:
+                    block = ctx().blocks.PreviewOption
+                case Mode.TUTORIAL:
+                    block = None
+                case _:
+                    assert_unreachable()
+            if block is not None:
+                Num._from_place_(BlockPlace(block, self.index))._set_(Num._accept_(value))
+            else:
+                raise RuntimeError("Options in the current mode cannot be set and use the default value")
         raise AttributeError("Options are read-only")
 
 

sonolus/script/quad.py

@@ -28,6 +28,20 @@ class Quad(Record):
     br: Vec2
     """The bottom-right corner of the quad."""
 
+    @classmethod
+    def zero(cls) -> Quad:
+        """Return a quad with all corners set to (0, 0).
+
+        Returns:
+            A new quad with all corners at the origin.
+        """
+        return cls(
+            bl=Vec2.zero(),
+            tl=Vec2.zero(),
+            tr=Vec2.zero(),
+            br=Vec2.zero(),
+        )
+
     @classmethod
     def from_quad(cls, value: QuadLike, /) -> Quad:
         """Create a quad from a quad-like value."""
@@ -80,7 +94,14 @@ class Quad(Record):
         ).translate(self.center * (Vec2(1, 1) - factor))
 
     def rotate(self, angle: float, /) -> Self:
-        """Rotate the quad by the given angle about the origin and return a new quad."""
+        """Rotate the quad by the given angle about the origin and return a new quad.
+
+        Args:
+            angle: The angle of rotation in radians. Positive angles rotate counterclockwise.
+
+        Returns:
+            A new quad rotated by the given angle.
+        """
         return Quad(
             bl=self.bl.rotate(angle),
             tl=self.tl.rotate(angle),
@@ -94,7 +115,15 @@ class Quad(Record):
         /,
         pivot: Vec2,
     ) -> Self:
-        """Rotate the quad by the given angle about the given pivot and return a new quad."""
+        """Rotate the quad by the given angle about the given pivot and return a new quad.
+
+        Args:
+            angle: The angle of rotation in radians. Positive angles rotate counterclockwise.
+            pivot: The pivot point for rotation.
+
+        Returns:
+            A new quad rotated about the pivot by the given angle.
+        """
         return Quad(
             bl=self.bl.rotate_about(angle, pivot),
             tl=self.tl.rotate_about(angle, pivot),
@@ -103,7 +132,14 @@ class Quad(Record):
         )
 
     def rotate_centered(self, angle: float, /) -> Self:
-        """Rotate the quad by the given angle about its center and return a new quad."""
+        """Rotate the quad by the given angle about its center and return a new quad.
+
+        Args:
+            angle: The angle of rotation in radians. Positive angles rotate counterclockwise.
+
+        Returns:
+            A new quad rotated about its center by the given angle.
+        """
         return self.rotate_about(angle, self.center)
 
     def permute(self, count: int = 1, /) -> Self:
@@ -147,6 +183,8 @@ class Quad(Record):
     def contains_point(self, point: Vec2, /) -> bool:
         """Check if the quad contains the given point.
 
+        It is not guaranteed whether points on the edges of the quad are considered inside or outside.
+
         Args:
             point: The point to check.
 

sonolus/script/record.py

@@ -20,8 +20,15 @@ from sonolus.script.internal.value import BackingSource, DataValue, Value
 from sonolus.script.num import Num
 
 
+class RecordMeta(type):
+    @meta_fn
+    def __pos__[T](cls: type[T]) -> T:
+        """Create a zero-initialized record instance."""
+        return cls._zero_()
+
+
 @dataclass_transform(eq_default=True)
-class Record(GenericValue):
+class Record(GenericValue, metaclass=RecordMeta):
     """Base class for user-defined data structures.
 
     Usage:
@@ -38,6 +45,15 @@ class Record(GenericValue):
             field1: T
             field2: U
         ```
+
+        Creating an instance:
+        ```python
+        record = MyRecord(field1=42, field2=True)
+        record_2 = MyGenericRecord[int, int](field1=42, field2=100)
+        record_3 = MyGenericRecord(field1=42, field2=100)  # Type arguments can be inferred
+        record_4 = +MyRecord  # Create a zero-initialized record
+        record_5 = +MyGenericRecord[int, int]
+        ```
     """
 
     _value: dict[str, Value]
@@ -239,12 +255,12 @@ class Record(GenericValue):
 
     def __str__(self):
         return (
-            f"{self.__class__.__name__}({", ".join(f"{field.name}={field.__get__(self)}" for field in self._fields)})"
+            f"{self.__class__.__name__}({', '.join(f'{field.name}={field.__get__(self)}' for field in self._fields)})"
         )
 
     def __repr__(self):
         return (
-            f"{self.__class__.__name__}({", ".join(f"{field.name}={field.__get__(self)!r}" for field in self._fields)})"
+            f"{self.__class__.__name__}({', '.join(f'{field.name}={field.__get__(self)!r}' for field in self._fields)})"
         )
 
     @meta_fn
@@ -268,6 +284,11 @@ class Record(GenericValue):
     def __hash__(self):
         return hash(tuple(field.__get__(self) for field in self._fields))
 
+    @meta_fn
+    def __pos__(self) -> Self:
+        """Return a copy of the record."""
+        return self._copy_()
+
     @classmethod
     @meta_fn
     def type_var_value(cls, var: TypeVar, /) -> Any: