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

sonolus/script/iterator.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from abc import abstractmethod
 from typing import Any
 
 from sonolus.script.internal.impl import meta_fn
+from sonolus.script.maybe import Maybe, Nothing, Some
 from sonolus.script.record import Record
 
 
@@ -12,57 +12,28 @@ class SonolusIterator[T]:
 
     This class is used to define custom iterators that can be used in Sonolus.py.
 
-    Inheritors must implement the `has_next`, `get`, and `advance` methods.
-    The `__next__` and `__iter__` methods are implemented by default.
+    Inheritors must implement the `next` method, which should return a `Maybe[T]`.
 
     Usage:
         ```python
         class MyIterator(Record, SonolusIterator):
-            def has_next(self) -> bool:
-                ...
-
-            def get(self) -> Any:
-                ...
-
-            def advance(self):
+            def next(self) -> Maybe[T]:
                 ...
         ```
     """
 
     _allow_instance_check_ = True
 
-    def next(self) -> T:
-        result = self.get()
-        self.advance()
-        return result
-
-    @abstractmethod
-    def has_next(self) -> bool:
-        """Return whether the iterator has more elements."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def get(self) -> T:
-        """Return the next element of the iterator.
-
-        May be called multiple times before calling `advance`.
-
-        Must not be called if `has_next` returns `False`.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def advance(self):
-        """Advance the iterator to the next element.
-
-        Must not be called if `has_next` returns `False`.
-        """
+    @meta_fn
+    def next(self) -> Maybe[T]:
         raise NotImplementedError
 
     def __next__(self) -> T:
-        if not self.has_next():
+        result = self.next()
+        if result.is_some:
+            return result.get_unsafe()
+        else:
             raise StopIteration
-        return self.next()
 
     def __iter__(self) -> SonolusIterator[T]:
         return self
@@ -73,15 +44,13 @@ class _Enumerator[V: SonolusIterator](Record, SonolusIterator):
     offset: int
     iterator: V
 
-    def has_next(self) -> bool:
-        return self.iterator.has_next()
-
-    def get(self) -> tuple[int, Any]:
-        return self.i + self.offset, self.iterator.get()
-
-    def advance(self):
+    def next(self) -> Maybe[tuple[int, Any]]:
+        value = self.iterator.next()
+        if value.is_nothing:
+            return Nothing
+        result = (self.i + self.offset, value.get_unsafe())
         self.i += 1
-        self.iterator.advance()
+        return Some(result)
 
 
 class _Zipper[T](Record, SonolusIterator):
@@ -89,11 +58,6 @@ class _Zipper[T](Record, SonolusIterator):
     iterators: T
 
     @meta_fn
-    def has_next(self) -> bool:
-        from sonolus.backend.visitor import compile_and_call
-
-        return compile_and_call(self._has_next, self._get_iterators())
-
     def _get_iterators(self) -> tuple[SonolusIterator, ...]:
         from sonolus.script.containers import Pair
 
@@ -105,81 +69,48 @@ class _Zipper[T](Record, SonolusIterator):
         iterators.append(v)
         return tuple(iterators)
 
-    def _has_next(self, iterators: tuple[SonolusIterator, ...]) -> bool:
-        for iterator in iterators:  # noqa: SIM110
-            if not iterator.has_next():
-                return False
-        return True
-
     @meta_fn
-    def get(self) -> tuple[Any, ...]:
+    def _get_next_values(self) -> tuple[Any, ...]:
         from sonolus.backend.visitor import compile_and_call
 
-        return tuple(compile_and_call(iterator.get) for iterator in self._get_iterators())
+        return tuple(compile_and_call(iterator.next) for iterator in self._get_iterators())
 
     @meta_fn
-    def advance(self):
+    def _values_to_tuple(self, values: tuple[Any, ...]) -> tuple[Any, ...]:
         from sonolus.backend.visitor import compile_and_call
 
-        for iterator in self._get_iterators():
-            compile_and_call(iterator.advance)
-
+        return tuple(compile_and_call(value.get_unsafe) for value in values)
 
-class _EmptyIterator(Record, SonolusIterator):
-    def has_next(self) -> bool:
-        return False
+    def next(self) -> Maybe[tuple[Any, ...]]:
+        values = self._get_next_values()
+        for value in values:
+            if value.is_nothing:
+                return Nothing
+        return Some(self._values_to_tuple(values))
 
-    def get(self) -> Any:
-        return None
 
-    def advance(self):
-        pass
+class _EmptyIterator(Record, SonolusIterator):
+    def next(self) -> Maybe[Any]:
+        return Nothing
 
 
 class _MappingIterator[T, Fn](Record, SonolusIterator):
     fn: Fn
     iterator: T
 
-    def has_next(self) -> bool:
-        return self.iterator.has_next()
-
-    def get(self) -> Any:
-        return self.fn(self.iterator.get())
-
-    def advance(self):
-        self.iterator.advance()
+    def next(self) -> Maybe[Any]:
+        return self.iterator.next().map(self.fn)
 
 
 class _FilteringIterator[T, Fn](Record, SonolusIterator):
     fn: Fn
     iterator: T
 
-    def has_next(self) -> bool:
-        while self.iterator.has_next():
-            if self.fn(self.iterator.get()):
-                return True
-            self.iterator.advance()
-        return False
-
-    def get(self) -> Any:
-        return self.iterator.get()
-
-    def advance(self):
-        self.iterator.advance()
-
-
-class _ChainingIterator[T](Record, SonolusIterator):
-    iterator: T
-
-    def has_next(self) -> bool:
-        return self.iterator.has_next()
-
-    def get(self) -> Any:
-        return self.iterator.get().get()
-
-    def advance(self):
-        self.iterator.get().advance()
-        while not self.iterator.get().has_next():
-            self.iterator.advance()
-            if not self.iterator.has_next():
-                break
+    def next(self) -> Maybe[T]:
+        while True:
+            value = self.iterator.next()
+            if value.is_nothing:
+                return Nothing
+            inside = value.get_unsafe()
+            if self.fn(inside):
+                return Some(inside)

sonolus/script/maybe.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from sonolus.script.internal.context import ctx
+from sonolus.script.internal.impl import meta_fn
+from sonolus.script.internal.transient import TransientValue
+from sonolus.script.internal.value import Value
+from sonolus.script.num import Num
+
+
+class Maybe[T](TransientValue):
+    """A type that either has a value or is empty.
+
+    Maybe has special behavior when returned from a function: it may be returned from multiple places
+    in a function, provided that all but one return statement returns the literal `Nothing`.
+
+    This type is not intended for other uses, such as being stored in a Record, Array, or Archetype.
+
+    Usage:
+        ```python
+        def fn(a, b):
+            if a:
+                return Some(b)
+            else:
+                return Nothing
+
+        result = fn(..., ...)
+        if result.is_some:
+            value = result.get()
+            ...
+        ```
+    """
+
+    _present: Num
+    _value: T
+
+    def __init__(self, *, present: bool, value: T):
+        self._present = Num._accept_(present)
+        self._value = value
+
+    @property
+    @meta_fn
+    def is_some(self) -> bool:
+        """Check if the value is present."""
+        if ctx():
+            return self._present._get_readonly_()
+        else:
+            return self._present._as_py_()
+
+    @property
+    def is_nothing(self) -> bool:
+        """Check if the value is empty."""
+        return not self.is_some
+
+    def get(self) -> T:
+        """Get the value if present, otherwise raise an error."""
+        assert self.is_some
+        return self.get_unsafe()
+
+    @meta_fn
+    def get_unsafe(self) -> T:
+        return self._value
+
+    def map[R](self, fn: Callable[[T], R], /) -> Maybe[R]:
+        if self.is_some:
+            return Some(fn(self.get_unsafe()))
+        return Nothing
+
+    @classmethod
+    def _accepts_(cls, value: Any) -> bool:
+        return isinstance(value, cls)
+
+    @classmethod
+    def _accept_(cls, value: Any) -> Maybe[T]:
+        if not cls._accepts_(value):
+            raise TypeError(f"Cannot accept value of type {type(value).__name__} as {cls.__name__}.")
+        return value
+
+    def _is_py_(self) -> bool:
+        return not self._present or not isinstance(self._value, Value) or self._value._is_py_()
+
+    def _as_py_(self) -> Any:
+        if not self._is_py_():
+            raise ValueError("Not a python value")
+        return self
+
+    def _copy_from_(self, value: Any):
+        raise TypeError("Maybe does not support mutation.")
+
+    def _copy_(self) -> Maybe[T]:
+        raise TypeError("Maybe does not support copying.")
+
+    def _set_(self, value: Any):
+        if not self._accepts_(value):
+            raise TypeError(f"Cannot set value of type {type(value).__name__} to {self.__class__.__name__}.")
+        if value is not Nothing and self._value is not value._value:
+            raise TypeError(f"Cannot set value of type {type(value._value).__name__} to {self.__class__.__name__}.")
+        self._present._set_(value._present)
+
+    @classmethod
+    def _get_merge_target_(cls, values: list[Any]) -> Any:
+        if not all(isinstance(v, cls) for v in values):
+            return NotImplemented
+        distinct = []
+        seen_ids = set()
+        for v in values:
+            if v is Nothing:
+                continue
+            if id(v._value) not in seen_ids:
+                distinct.append(v)
+                seen_ids.add(id(v._value))
+        match distinct:
+            case []:
+                return Nothing
+            case [v]:
+                return Maybe(present=Num._alloc_(), value=v._value)
+            case _:
+                return NotImplemented
+
+
+def Some[T](value: T) -> Maybe[T]:  # noqa: N802
+    """Create a `Maybe` instance with a value.
+
+    Args:
+        value: The contained value.
+
+    Returns:
+        A `Maybe` instance that contains the provided value.
+    """
+    return Maybe(present=True, value=value)
+
+
+Nothing: Maybe[Any] = Maybe(present=False, value=None)  # type: ignore
+
+# Note: has to come after the definition to hide the definition in the docs.
+Nothing: Maybe[Any]
+"""The empty `Maybe` instance."""

sonolus/script/num.py

@@ -1,9 +1,11 @@
+# type: ignore
 from __future__ import annotations
 
 import operator
 from collections.abc import Callable, Iterable
 from typing import TYPE_CHECKING, Any, Self, TypeGuard, final, runtime_checkable
 
+from sonolus.backend.blocks import BlockData
 from sonolus.backend.ir import IRConst, IRExpr, IRGet, IRPureInstr, IRSet
 from sonolus.backend.ops import Op
 from sonolus.backend.place import BlockPlace
@@ -114,15 +116,28 @@ class _Num(Value, metaclass=_NumMeta):
 
     def _get_(self) -> Self:
         if ctx():
+            if isinstance(self.data, BlockPlace):
+                ctx().check_readable(self.data)
             place = ctx().alloc(size=1)
+            ctx().add_statements(IRSet(place, self.ir()))
+            return Num(place)
+        else:
+            return Num(self.data)
+
+    def _get_readonly_(self) -> Self:
+        if ctx():
             if isinstance(self.data, BlockPlace):
                 ctx().check_readable(self.data)
+                if isinstance(self.data.block, BlockData) and not ctx().is_writable(self.data):
+                    # This block is immutable in the current callback, so no need to copy it in case it changes.
+                    return Num(self.data)
+            place = ctx().alloc(size=1)
             ctx().add_statements(IRSet(place, self.ir()))
             return Num(place)
         else:
             return Num(self.data)
 
-    def _set_(self, value: Self):
+    def _set_(self, value: Any):
         value = Num._accept_(value)
         if ctx():
             match self.data:
@@ -136,7 +151,7 @@ class _Num(Value, metaclass=_NumMeta):
         else:
             self.data = value.data
 
-    def _copy_from_(self, value: Self):
+    def _copy_from_(self, value: Any):
         raise ValueError("Cannot assign to a number")
 
     def _copy_(self) -> Self:
@@ -439,17 +454,17 @@ if TYPE_CHECKING:
 
     @runtime_checkable
     class Num[T](Protocol, int, bool, float):
-        def __add__(self, other: T, /) -> Num: ...
-        def __sub__(self, other: T, /) -> Num: ...
-        def __mul__(self, other: T, /) -> Num: ...
-        def __truediv__(self, other: T, /) -> Num: ...
-        def __floordiv__(self, other: T, /) -> Num: ...
-        def __mod__(self, other: T, /) -> Num: ...
-        def __pow__(self, other: T, /) -> Num: ...
-
-        def __neg__(self, /) -> Num: ...
-        def __pos__(self, /) -> Num: ...
-        def __abs__(self, /) -> Num: ...
+        def __add__(self, other: T, /) -> Num | int | bool | float: ...
+        def __sub__(self, other: T, /) -> Num | int | bool | float: ...
+        def __mul__(self, other: T, /) -> Num | int | bool | float: ...
+        def __truediv__(self, other: T, /) -> Num | int | bool | float: ...
+        def __floordiv__(self, other: T, /) -> Num | int | bool | float: ...
+        def __mod__(self, other: T, /) -> Num | int | bool | float: ...
+        def __pow__(self, other: T, /) -> Num | int | bool | float: ...
+
+        def __neg__(self, /) -> Num | int | bool | float: ...
+        def __pos__(self, /) -> Num | int | bool | float: ...
+        def __abs__(self, /) -> Num | int | bool | float: ...
 
         def __eq__(self, other: Any, /) -> bool: ...
         def __ne__(self, other: Any, /) -> bool: ...
@@ -461,10 +476,10 @@ if TYPE_CHECKING:
         def __hash__(self, /) -> int: ...
 
         @property
-        def real(self) -> Num: ...
+        def real(self) -> Num | int | bool | float: ...
 
         @property
-        def imag(self) -> Num: ...
+        def imag(self) -> Num | int | bool | float: ...
 else:
     # Need to do this to satisfy type checkers (especially Pycharm)
     _Num.__name__ = "Num"

sonolus/script/options.py

@@ -175,7 +175,7 @@ def select_option(
     return _SelectOption(name, description, standard, advanced, scope, default, values)
 
 
-type Options = NewType("Options", Any)
+type Options = NewType("Options", Any)  # type: ignore
 type _OptionInfo = _SliderOption | _ToggleOption | _SelectOption
 
 

sonolus/script/particle.py

@@ -95,7 +95,7 @@ def particle(name: str) -> Any:
     return _ParticleInfo(name)
 
 
-type Particles = NewType("Particles", Any)
+type Particles = NewType("Particles", Any)  # type: ignore
 
 
 @dataclass_transform()

sonolus/script/pointer.py

@@ -6,7 +6,7 @@ from sonolus.script.num import Num, _is_num
 
 
 @meta_fn
-def _deref[T: Value](block: Num, offset: Num, type_: type[T]) -> T:
+def _deref[T: Value](block: int, offset: int, type_: type[T]) -> T:
     block = Num._accept_(block)
     offset = Num._accept_(offset)
     type_ = validate_value(type_)._as_py_()

sonolus/script/project.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
 from os import PathLike
 from pathlib import Path
-from typing import ClassVar, Self, TypedDict
+from typing import ClassVar, TypedDict
 
 from sonolus.backend.optimize import optimize
 from sonolus.backend.optimize.passes import CompilerPass
@@ -25,14 +25,23 @@ class Project:
     def __init__(
         self,
         engine: Engine,
-        levels: list[Level] | None = None,
+        levels: Iterable[Level] | Callable[[], Iterable[Level]] | None = None,
         resources: PathLike | None = None,
     ):
         self.engine = engine
-        self.levels = levels or []
+        match levels:
+            case Callable():
+                self._level_source = lazy_loader(levels)
+            case Iterable():
+                self._level_source = levels
+            case None:
+                self._level_source = []
+            case _:
+                raise TypeError(f"Invalid type for levels: {type(levels)}. Expected Iterable or Callable.")
+        self._levels = None
         self.resources = Path(resources or "resources")
 
-    def with_levels(self, levels: list[Level]) -> Self:
+    def with_levels(self, levels: list[Level]) -> Project:
         """Create a new project with the specified levels.
 
         Args:
@@ -78,6 +87,16 @@ class Project:
 
         return get_project_schema(self)
 
+    @property
+    def levels(self) -> list[Level]:
+        if self._levels is None:
+            self._levels = list(self._level_source)
+        return self._levels
+
+
+def lazy_loader(fn):
+    yield from fn()
+
 
 class ProjectSchema(TypedDict):
     archetypes: list[ArchetypeSchema]

sonolus/script/quad.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Protocol, Self
+from typing import Protocol
 
 from sonolus.script.record import Record
 from sonolus.script.values import zeros
@@ -57,7 +57,7 @@ class Quad(Record):
         """The center of the quad."""
         return (self.bl + self.tr + self.tl + self.br) / 4
 
-    def translate(self, translation: Vec2, /) -> Self:
+    def translate(self, translation: Vec2, /) -> Quad:
         """Translate the quad by the given translation and return a new quad."""
         return Quad(
             bl=self.bl + translation,
@@ -66,7 +66,7 @@ class Quad(Record):
             br=self.br + translation,
         )
 
-    def scale(self, factor: Vec2, /) -> Self:
+    def scale(self, factor: Vec2, /) -> Quad:
         """Scale the quad by the given factor about the origin and return a new quad."""
         return Quad(
             bl=self.bl * factor,
@@ -75,7 +75,7 @@ class Quad(Record):
             br=self.br * factor,
         )
 
-    def scale_about(self, factor: Vec2, /, pivot: Vec2) -> Self:
+    def scale_about(self, factor: Vec2, /, pivot: Vec2) -> Quad:
         """Scale the quad by the given factor about the given pivot and return a new quad."""
         return Quad(
             bl=(self.bl - pivot) * factor + pivot,
@@ -84,7 +84,7 @@ class Quad(Record):
             br=(self.br - pivot) * factor + pivot,
         )
 
-    def scale_centered(self, factor: Vec2, /) -> Self:
+    def scale_centered(self, factor: Vec2, /) -> Quad:
         """Scale the quad by the given factor about its center and return a new quad."""
         return Quad(
             bl=self.bl * factor,
@@ -93,7 +93,7 @@ class Quad(Record):
             br=self.br * factor,
         ).translate(self.center * (Vec2(1, 1) - factor))
 
-    def rotate(self, angle: float, /) -> Self:
+    def rotate(self, angle: float, /) -> Quad:
         """Rotate the quad by the given angle about the origin and return a new quad.
 
         Args:
@@ -114,7 +114,7 @@ class Quad(Record):
         angle: float,
         /,
         pivot: Vec2,
-    ) -> Self:
+    ) -> Quad:
         """Rotate the quad by the given angle about the given pivot and return a new quad.
 
         Args:
@@ -131,7 +131,7 @@ class Quad(Record):
             br=self.br.rotate_about(angle, pivot),
         )
 
-    def rotate_centered(self, angle: float, /) -> Self:
+    def rotate_centered(self, angle: float, /) -> Quad:
         """Rotate the quad by the given angle about its center and return a new quad.
 
         Args:
@@ -142,7 +142,7 @@ class Quad(Record):
         """
         return self.rotate_about(angle, self.center)
 
-    def permute(self, count: int = 1, /) -> Self:
+    def permute(self, count: int = 1, /) -> Quad:
         """Perform a cyclic permutation of the quad's vertices and return a new quad.
 
         On a square, this operation is equivalent to rotating the square counterclockwise 90 degrees `count` times.
@@ -269,7 +269,7 @@ class Rect(Record):
             br=self.br,
         )
 
-    def translate(self, translation: Vec2, /) -> Self:
+    def translate(self, translation: Vec2, /) -> Rect:
         """Translate the rectangle by the given translation and return a new rectangle."""
         return Rect(
             t=self.t + translation.y,
@@ -278,7 +278,7 @@ class Rect(Record):
             l=self.l + translation.x,
         )
 
-    def scale(self, factor: Vec2, /) -> Self:
+    def scale(self, factor: Vec2, /) -> Rect:
         """Scale the rectangle by the given factor about the origin and return a new rectangle."""
         return Rect(
             t=self.t * factor.y,
@@ -287,7 +287,7 @@ class Rect(Record):
             l=self.l * factor.x,
         )
 
-    def scale_about(self, factor: Vec2, /, pivot: Vec2) -> Self:
+    def scale_about(self, factor: Vec2, /, pivot: Vec2) -> Rect:
         """Scale the rectangle by the given factor about the given pivot and return a new rectangle."""
         return Rect(
             t=(self.t - pivot.y) * factor.y + pivot.y,
@@ -296,7 +296,7 @@ class Rect(Record):
             l=(self.l - pivot.x) * factor.x + pivot.x,
         )
 
-    def scale_centered(self, factor: Vec2, /) -> Self:
+    def scale_centered(self, factor: Vec2, /) -> Rect:
         """Scale the rectangle by the given factor about its center and return a new rectangle."""
         return Rect(
             t=self.t * factor.y,
@@ -305,7 +305,7 @@ class Rect(Record):
             l=self.l * factor.x,
         ).translate(self.center * (Vec2(1, 1) - factor))
 
-    def expand(self, expansion: Vec2, /) -> Self:
+    def expand(self, expansion: Vec2, /) -> Rect:
         """Expand the rectangle by the given amount and return a new rectangle."""
         return Rect(
             t=self.t + expansion.y,
@@ -314,7 +314,7 @@ class Rect(Record):
             l=self.l - expansion.x,
         )
 
-    def shrink(self, shrinkage: Vec2, /) -> Self:
+    def shrink(self, shrinkage: Vec2, /) -> Rect:
         """Shrink the rectangle by the given amount and return a new rectangle."""
         return Rect(
             t=self.t - shrinkage.y,