sonolus.py 0.3.4__py3-none-any.whl → 0.4.1__py3-none-any.whl

sonolus/script/iterator.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
-from abc import abstractmethod
+from collections.abc import Iterator
 from typing import Any
 
+from sonolus.script.internal.context import ctx
 from sonolus.script.internal.impl import meta_fn
+from sonolus.script.maybe import Maybe, Nothing, Some
 from sonolus.script.record import Record
 
 
@@ -12,57 +14,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 +46,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 +60,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 +71,61 @@ 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)
+
+
+@meta_fn
+def maybe_next[T](iterator: Iterator[T]) -> Maybe[T]:
+    """Get the next item from an iterator as a `Maybe` if it exists or `Nothing` otherwise."""
+    from sonolus.backend.visitor import compile_and_call
+
+    if not isinstance(iterator, SonolusIterator):
+        raise TypeError("Iterator must be an instance of SonolusIterator.")
+    if ctx():
+        return compile_and_call(iterator.next)
+    else:
+        return iterator.next()

sonolus/script/level.py

@@ -145,7 +145,7 @@ class Level:
         )
 
 
-type EntityListArg = list[PlayArchetype | EntityListArg]
+type EntityListArg = list[list[PlayArchetype] | PlayArchetype] | PlayArchetype
 
 
 def flatten_entities(entities: EntityListArg) -> Iterator[PlayArchetype]:
@@ -175,7 +175,7 @@ class LevelData:
     bgm_offset: float
     entities: list[PlayArchetype]
 
-    def __init__(self, bgm_offset: float, entities: list[PlayArchetype]) -> None:
+    def __init__(self, bgm_offset: float, entities: EntityListArg) -> None:
         self.bgm_offset = bgm_offset
         self.entities = [*flatten_entities(entities)]
 

sonolus/script/maybe.py

@@ -0,0 +1,241 @@
+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, validate_value
+from sonolus.script.internal.transient import TransientValue
+from sonolus.script.internal.value import Value
+from sonolus.script.num import Num
+from sonolus.script.values import copy, zeros
+
+
+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 = validate_value(value)
+
+    @property
+    @meta_fn
+    def is_some(self) -> bool:
+        """Check if the value is present."""
+        if ctx():
+            if self._present._is_py_():
+                # Makes this a compile time constant.
+                return self._present
+            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:
+        if ctx():
+            return self._value
+        else:
+            return self._value._as_py_()
+
+    def map[R](self, fn: Callable[[T], R], /) -> Maybe[R]:
+        """Map the contained value to a new value using the provided function.
+
+        If the value is not present, returns `Nothing`.
+
+        Args:
+            fn: A function that takes the contained value and returns a new value.
+
+        Returns:
+            A `Maybe` instance containing the result of the function if the value is present, otherwise `Nothing`.
+        """
+        if self.is_some:
+            return Some(fn(self.get_unsafe()))
+        return Nothing
+
+    def flat_map[R](self, fn: Callable[[T], Maybe[R]], /) -> Maybe[R]:
+        """Flat map the contained value to a new `Maybe` using the provided function.
+
+        If the value is not present, returns `Nothing`.
+
+        Args:
+            fn: A function that takes the contained value and returns a new `Maybe`.
+
+        Returns:
+            A `Maybe` instance containing the result of the function if the value is present, otherwise `Nothing`.
+        """
+        if self.is_some:
+            return fn(self.get_unsafe())
+        return Nothing
+
+    def or_default(self, default: T) -> T:
+        """Return a copy of the contained value if present, otherwise return a copy of the given default value.
+
+        Args:
+            default: The default value to return if the contained value is not present.
+
+        Returns:
+            A copy of the contained value if present, otherwise a copy of the default value.
+        """
+        result = _box(copy(default))
+        if self.is_some:
+            result.value = self.get_unsafe()
+        return result.value
+
+    @meta_fn
+    def or_else(self, fn: Callable[[], T], /) -> T:
+        """Return a copy of the contained value if present, otherwise return a copy of the result of the given function.
+
+        Args:
+            fn: A function that returns a value to use if the contained value is not present.
+
+        Returns:
+            A copy of the contained value if present, otherwise a copy of the result of calling the function.
+        """
+        from sonolus.backend.visitor import compile_and_call
+
+        if ctx():
+            if self.is_some._is_py_():  # type: ignore
+                if self.is_some._as_py_():  # type: ignore
+                    return copy(self.get_unsafe())
+                else:
+                    return copy(compile_and_call(fn))
+            else:
+                return compile_and_call(self._or_else, fn)
+        elif self.is_some:
+            return copy(self.get_unsafe())
+        else:
+            return copy(fn())
+
+    def _or_else(self, fn: Callable[[], T], /) -> T:
+        result = _box(zeros(self.contained_type))
+        if self.is_some:
+            result.value = self.get_unsafe()
+        else:
+            result.value = fn()
+        return result.value
+
+    @property
+    def tuple(self) -> tuple[bool, T]:
+        """Return whether the value is present and a copy of the contained value if present as a tuple.
+
+        If the value is not present, the tuple will contain `False` and a zero initialized value of the contained type.
+        """
+        result_value = _box(zeros(self.contained_type))
+        if self.is_some:
+            result_value.value = self.get_unsafe()
+        return self.is_some, result_value.value
+
+    @property
+    @meta_fn
+    def contained_type(self):
+        return type(self._value)
+
+    @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."""
+
+
+@meta_fn
+def _box(value):
+    from sonolus.script.containers import Box
+
+    return Box(value)

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:
@@ -438,14 +453,14 @@ if TYPE_CHECKING:
     from typing import Protocol
 
     @runtime_checkable
-    class Num[T](Protocol, int, bool, float):
-        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: ...
+    class Num(Protocol, int, bool, float):
+        def __add__(self, other: Any, /) -> Num | int | bool | float: ...
+        def __sub__(self, other: Any, /) -> Num | int | bool | float: ...
+        def __mul__(self, other: Any, /) -> Num | int | bool | float: ...
+        def __truediv__(self, other: Any, /) -> Num | int | bool | float: ...
+        def __floordiv__(self, other: Any, /) -> Num | int | bool | float: ...
+        def __mod__(self, other: Any, /) -> Num | int | bool | float: ...
+        def __pow__(self, other: Any, /) -> Num | int | bool | float: ...
 
         def __neg__(self, /) -> Num | int | bool | float: ...
         def __pos__(self, /) -> Num | int | bool | float: ...
@@ -453,10 +468,10 @@ if TYPE_CHECKING:
 
         def __eq__(self, other: Any, /) -> bool: ...
         def __ne__(self, other: Any, /) -> bool: ...
-        def __lt__(self, other: T, /) -> bool: ...
-        def __le__(self, other: T, /) -> bool: ...
-        def __gt__(self, other: T, /) -> bool: ...
-        def __ge__(self, other: T, /) -> bool: ...
+        def __lt__(self, other: Any, /) -> bool: ...
+        def __le__(self, other: Any, /) -> bool: ...
+        def __gt__(self, other: Any, /) -> bool: ...
+        def __ge__(self, other: Any, /) -> bool: ...
 
         def __hash__(self, /) -> int: ...
 

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/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: Iterable[Level] | Callable[[], Iterable[Level]] | None) -> 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]