add return types to functions

lenskit · May 23, 2024 · d3e283c · d3e283c
1 parent ce16a56
commit d3e283c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/seedbank/__init__.py b/seedbank/__init__.py
@@ -11,6 +11,7 @@
 from types import ModuleType
 
 import numpy as np
+from numpy.random import SeedSequence
 from typing_extensions import Optional
 
 from seedbank._keys import RNGKey, SeedLike, make_seed
@@ -47,7 +48,7 @@
 ]
 
 
-def initialize(seed: SeedLike, *keys: RNGKey):
+def initialize(seed: SeedLike, *keys: RNGKey) -> SeedSequence:
     """
     Initialize the random infrastructure with a seed.  This function should generally be
     called very early in the setup.  This initializes all known and available RNGs with
@@ -78,7 +79,7 @@ def initialize(seed: SeedLike, *keys: RNGKey):
     return _root_state.seed
 
 
-def derive_seed(*keys: RNGKey, base: Optional[np.random.SeedSequence] = None):
+def derive_seed(*keys: RNGKey, base: Optional[np.random.SeedSequence] = None) -> SeedSequence:
     """
     Derive a seed from the root seed, optionally with additional seed keys.
 
@@ -96,7 +97,7 @@ def derive_seed(*keys: RNGKey, base: Optional[np.random.SeedSequence] = None):
     return _root_state.derive(base, keys).seed
 
 
-def root_seed():
+def root_seed() -> SeedSequence:
     """
     Get the current root seed.
 

diff --git a/seedbank/cupy.py b/seedbank/cupy.py
@@ -29,7 +29,7 @@ def seed(state):
     cupy.random.seed(state.int_seed)
 
 
-def cupy_rng(spec: Optional[SeedLike | cupy.random.Generator] = None):
+def cupy_rng(spec: Optional[SeedLike | cupy.random.Generator] = None) -> cupy.random.Generator:
     """
     Get a CuPy random number generator.  This works like :func:`numpy_rng`, but
     it returns a :class:`cupy.random.Generator` instead.