Random methods on resonate ctx (#209)

Author: Tomperez98

resonate/dependencies.py

@@ -10,5 +10,5 @@ def __init__(self) -> None:
     def add(self, key: str, obj: Any) -> None:
         self._deps[key] = obj
 
-    def get(self, key: str) -> Any:
-        return self._deps[key]
+    def get[T](self, key: str, default: T) -> Any | T:
+        return self._deps.get(key, default)

resonate/resonate.py

@@ -202,6 +202,32 @@ def promises(self) -> PromiseStore:
         return self._store.promises
 
 
+class Random:
+    def __init__(self, ctx: Context) -> None:
+        self.ctx = ctx
+
+    def random(self) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).random())
+
+    def betavariate(self, alpha: float, beta: float) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).betavariate(alpha, beta))
+
+    def randrange(self, start: int, stop: int | None = None, step: int = 1) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).randrange(start, stop, step))
+
+    def randint(self, a: int, b: int) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).randint(a, b))
+
+    def getrandbits(self, k: int) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).getrandbits(k))
+
+    def triangular(self, low: float = 0, high: float = 1, mode: float | None = None) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).triangular(low, high, mode))
+
+    def expovariate(self, lambd: float = 1) -> LFC:
+        return self.ctx.lfc(lambda _: self.ctx.get_dependency("resonate:random", random).expovariate(lambd))
+
+
 # Context
 class Context:
     def __init__(self, id: str, info: Info, opts: Options, registry: Registry, dependencies: Dependencies) -> None:
@@ -211,6 +237,7 @@ def __init__(self, id: str, info: Info, opts: Options, registry: Registry, depen
         self._registry = registry
         self._dependencies = dependencies
         self._counter = 0
+        self._random = Random(self)
 
     @property
     def id(self) -> str:
@@ -220,6 +247,10 @@ def id(self) -> str:
     def info(self) -> Info:
         return self._info
 
+    @property
+    def random(self) -> Random:
+        return self._random
+
     @overload
     def lfi[**P, R](self, func: Callable[Concatenate[Context, P], Generator[Any, Any, R]], *args: P.args, **kwargs: P.kwargs) -> LFI: ...
     @overload
@@ -285,11 +316,8 @@ def promise(self, data: Any = None, headers: dict[str, str] | None = None) -> RF
         self._counter += 1
         return RFI(f"{self.id}.{self._counter}", Base(data, headers))
 
-    def random(self, a: int, b: int) -> LFC:
-        return self.lfc(lambda _, a, b: random.randint(a, b), a, b)
-
-    def get_dependency(self, name: str) -> Any:
-        return self._dependencies.get(name)
+    def get_dependency[T](self, key: str, default: T = None) -> Any | T:
+        return self._dependencies.get(key, default)
 
     def _lfi_func(self, f: str | Callable) -> tuple[Callable, int, dict[int, Callable] | None]:
         match f:

tests/test_bridge.py

@@ -105,7 +105,9 @@ def add_one(ctx: Context, n: int) -> int:
 
 
 def get_dependency(ctx: Context) -> int:
-    return ctx.get_dependency("foo") + 1
+    dep = ctx.get_dependency("foo")
+    assert dep is not None
+    return dep + 1
 
 
 def rfi_add_one_by_name(ctx: Context, n: int) -> Generator[Any, Any, int]:
@@ -123,7 +125,7 @@ def hitl(ctx: Context, id: str | None) -> Generator[Yieldable, Any, int]:
 
 
 def random_generation(ctx: Context) -> Generator[Yieldable, Any, float]:
-    return (yield ctx.random(0, 10))
+    return (yield ctx.random.randint(0, 10))
 
 
 def get_stores_config() -> list[tuple[Store, MessageSource | None]]: