sonolus.py 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

sonolus/script/num.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import operator
 from collections.abc import Callable, Iterable
-from typing import TYPE_CHECKING, Any, Self, TypeIs, final
+from typing import TYPE_CHECKING, Any, Self, TypeIs, final, runtime_checkable
 
 from sonolus.backend.ir import IRConst, IRGet, IRPureInstr, IRSet
 from sonolus.backend.ops import Op
@@ -14,18 +14,18 @@ from sonolus.script.internal.impl import meta_fn
 from sonolus.script.internal.value import Value
 
 
-class NumMeta(type):
+class _NumMeta(type):
     def __instancecheck__(cls, instance):
-        return isinstance(instance, float | int | bool) or is_num(instance)
+        return isinstance(instance, float | int | bool) or _is_num(instance)
 
 
-def is_num(value: Any) -> TypeIs[Num]:
+def _is_num(value: Any) -> TypeIs[Num]:
     """Check if a value is a precisely Num instance."""
     return type.__instancecheck__(Num, value)  # type: ignore # noqa: PLC2801
 
 
 @final
-class _Num(Value, metaclass=NumMeta):
+class _Num(Value, metaclass=_NumMeta):
     # This works for ints, floats, and bools
     # Since we don't support complex numbers, real is equal to the original number
     __match_args__ = ("real",)
@@ -37,7 +37,7 @@ class _Num(Value, metaclass=NumMeta):
             raise TypeError("Cannot create a Num from a complex number")
         if isinstance(data, int):
             data = float(data)
-        if is_num(data):
+        if _is_num(data):
             raise InternalError("Cannot create a Num from a Num")
         self.data = data
 
@@ -73,7 +73,7 @@ class _Num(Value, metaclass=NumMeta):
     def _accept_(cls, value: Any) -> Self:
         if not cls._accepts_(value):
             raise TypeError(f"Cannot accept {value}")
-        if is_num(value):
+        if _is_num(value):
             return value
         return cls(value)
 
@@ -92,7 +92,7 @@ class _Num(Value, metaclass=NumMeta):
         value = next(iter(values))
         return Num(value)
 
-    def _to_list_(self) -> list[float | BlockPlace]:
+    def _to_list_(self, level_refs: dict[Any, int] | None = None) -> list[float | BlockPlace]:
         return [self.data]
 
     @classmethod
@@ -122,7 +122,10 @@ class _Num(Value, metaclass=NumMeta):
         raise ValueError("Cannot assign to a number")
 
     def _copy_(self) -> Self:
-        return self
+        if ctx():
+            return self._get_()
+        else:
+            return Num(self.data)
 
     @classmethod
     def _alloc_(cls) -> Self:
@@ -168,7 +171,7 @@ class _Num(Value, metaclass=NumMeta):
     def __eq__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data == b.data)
+                return Num(a._as_py_() == b._as_py_())
             if a.data == b.data:
                 return Num(True)
             return None
@@ -184,7 +187,7 @@ class _Num(Value, metaclass=NumMeta):
     def __ne__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data != b.data)
+                return Num(a._as_py_() != b._as_py_())
             if a.data == b.data:
                 return Num(False)
             return None
@@ -195,7 +198,7 @@ class _Num(Value, metaclass=NumMeta):
     def __lt__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data < b.data)
+                return Num(a._as_py_() < b._as_py_())
             if a.data == b.data:
                 return Num(False)
             return None
@@ -206,7 +209,9 @@ class _Num(Value, metaclass=NumMeta):
     def __le__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data <= b.data)
+                return Num(a._as_py_() <= b._as_py_())
+            if a.data == b.data:
+                return Num(True)
             return None
 
         return self._bin_op(other, const_fn, Op.LessOr)
@@ -215,7 +220,7 @@ class _Num(Value, metaclass=NumMeta):
     def __gt__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data > b.data)
+                return Num(a._as_py_() > b._as_py_())
             if a.data == b.data:
                 return Num(False)
             return None
@@ -226,7 +231,9 @@ class _Num(Value, metaclass=NumMeta):
     def __ge__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data >= b.data)
+                return Num(a._as_py_() >= b._as_py_())
+            if a.data == b.data:
+                return Num(True)
             return None
 
         return self._bin_op(other, const_fn, Op.GreaterOr)
@@ -235,10 +242,10 @@ class _Num(Value, metaclass=NumMeta):
     def __add__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data + b.data)
-            if a._is_py_() and a.data == 0:
+                return Num(a._as_py_() + b._as_py_())
+            if a._is_py_() and a._as_py_() == 0:
                 return b
-            if b._is_py_() and b.data == 0:
+            if b._is_py_() and b._as_py_() == 0:
                 return a
             return None
 
@@ -248,10 +255,10 @@ class _Num(Value, metaclass=NumMeta):
     def __sub__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data - b.data)
-            if a._is_py_() and a.data == 0:
+                return Num(a._as_py_() - b._as_py_())
+            if a._is_py_() and a._as_py_() == 0:
                 return -b
-            if b._is_py_() and b.data == 0:
+            if b._is_py_() and b._as_py_() == 0:
                 return a
             return None
 
@@ -261,16 +268,16 @@ class _Num(Value, metaclass=NumMeta):
     def __mul__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data * b.data)
+                return Num(a._as_py_() * b._as_py_())
             if a._is_py_():
-                if a.data == 0:
+                if a._as_py_() == 0:
                     return Num(0)
-                if a.data == 1:
+                if a._as_py_() == 1:
                     return b
             if b._is_py_():
-                if b.data == 0:
+                if b._as_py_() == 0:
                     return Num(0)
-                if b.data == 1:
+                if b._as_py_() == 1:
                     return a
             return None
 
@@ -280,13 +287,13 @@ class _Num(Value, metaclass=NumMeta):
     def __truediv__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                if b.data == 0:
+                if b._as_py_() == 0:
                     return None
-                return Num(a.data / b.data)
+                return Num(a._as_py_() / b._as_py_())
             if b._is_py_():
-                if b.data == 1:
+                if b._as_py_() == 1:
                     return a
-                if b.data == -1:
+                if b._as_py_() == -1:
                     return -a
             return None
 
@@ -296,13 +303,13 @@ class _Num(Value, metaclass=NumMeta):
     def __floordiv__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                if b.data == 0:
+                if b._as_py_() == 0:
                     return None
-                return Num(a.data // b.data)
+                return Num(a._as_py_() // b._as_py_())
             if b._is_py_():
-                if b.data == 1:
+                if b._as_py_() == 1:
                     return a
-                if b.data == -1:
+                if b._as_py_() == -1:
                     return -a
             return None
 
@@ -312,9 +319,9 @@ class _Num(Value, metaclass=NumMeta):
     def __mod__(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                if b.data == 0:
+                if b._as_py_() == 0:
                     return None
-                return Num(a.data % b.data)
+                return Num(a._as_py_() % b._as_py_())
             return None
 
         return self._bin_op(other, const_fn, Op.Mod)
@@ -324,13 +331,13 @@ class _Num(Value, metaclass=NumMeta):
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
                 try:
-                    return Num(a.data**b.data)
+                    return Num(a._as_py_() ** b._as_py_())
                 except OverflowError:
                     return None
             if b._is_py_():
-                if b.data == 0:
+                if b._as_py_() == 0:
                     return Num(1)
-                if b.data == 1:
+                if b._as_py_() == 1:
                     return a
             return None
 
@@ -360,9 +367,9 @@ class _Num(Value, metaclass=NumMeta):
     def and_(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data and b.data)
+                return Num(a._as_py_() and b._as_py_())
             if a._is_py_():
-                if a.data == 0:
+                if a._as_py_() == 0:
                     return a
                 else:
                     return b
@@ -373,9 +380,9 @@ class _Num(Value, metaclass=NumMeta):
     def or_(self, other) -> Self:
         def const_fn(a: Self, b: Self) -> Num | None:
             if a._is_py_() and b._is_py_():
-                return Num(a.data or b.data)
+                return Num(a._as_py_() or b._as_py_())
             if a._is_py_():
-                if a.data == 0:
+                if a._as_py_() == 0:
                     return b
                 else:
                     return a
@@ -396,11 +403,36 @@ class _Num(Value, metaclass=NumMeta):
 
 
 if TYPE_CHECKING:
-
-    class __Num(float, int, bool, _Num):  # type: ignore
-        pass
-
-    Num = __Num | _Num | float | int | bool
+    from typing import Protocol
+
+    @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 __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 __hash__(self, /) -> int: ...
+
+        @property
+        def real(self) -> Num: ...
+
+        @property
+        def imag(self) -> Num: ...
 else:
     # Need to do this to satisfy type checkers (especially Pycharm)
     _Num.__name__ = "Num"