docs: 🏷️ improve type hints for internal TorchRandomWrapper object.

qthequartermasterman · May 14, 2024 · af953ce · af953ce
1 parent 12ec834
commit af953ce
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/hypothesis_torch/register_random_torch_state.py b/hypothesis_torch/register_random_torch_state.py
@@ -1,5 +1,7 @@
 """Register a random.Random-compatible shim for `torch.random` with Hypothesis."""
 
+from typing import Any, Callable
+
 import hypothesis
 import hypothesis.internal.entropy
 import torch
@@ -11,6 +13,10 @@ class TorchRandomWrapper:
     Hypothesis's random state manager requires a random.Random-compatible object.
     """
 
+    seed: Callable[..., Any]
+    getstate: Callable[[], Any]
+    setstate: Callable[..., Any]
+
     def __init__(self) -> None:
         """Initialize the TorchRandomWrapper."""
         self.seed = torch.random.manual_seed