bit_generator.pyi 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. import abc
  2. from _typeshed import Incomplete
  3. from collections.abc import Callable, Mapping, Sequence
  4. from threading import Lock
  5. from typing import (
  6. Any,
  7. ClassVar,
  8. Literal,
  9. NamedTuple,
  10. Self,
  11. TypeAlias,
  12. TypedDict,
  13. overload,
  14. type_check_only,
  15. )
  16. from typing_extensions import CapsuleType
  17. import numpy as np
  18. from numpy._typing import (
  19. NDArray,
  20. _ArrayLikeInt_co,
  21. _DTypeLike,
  22. _ShapeLike,
  23. _UInt32Codes,
  24. _UInt64Codes,
  25. )
  26. __all__ = ["BitGenerator", "SeedSequence"]
  27. ###
  28. _DTypeLikeUint_: TypeAlias = _DTypeLike[np.uint32 | np.uint64] | _UInt32Codes | _UInt64Codes
  29. @type_check_only
  30. class _SeedSeqState(TypedDict):
  31. entropy: int | Sequence[int] | None
  32. spawn_key: tuple[int, ...]
  33. pool_size: int
  34. n_children_spawned: int
  35. @type_check_only
  36. class _Interface(NamedTuple):
  37. state_address: Incomplete
  38. state: Incomplete
  39. next_uint64: Incomplete
  40. next_uint32: Incomplete
  41. next_double: Incomplete
  42. bit_generator: Incomplete
  43. @type_check_only
  44. class _CythonMixin:
  45. def __setstate_cython__(self, pyx_state: object, /) -> None: ...
  46. def __reduce_cython__(self) -> Any: ... # noqa: ANN401
  47. @type_check_only
  48. class _GenerateStateMixin(_CythonMixin):
  49. def generate_state(self, /, n_words: int, dtype: _DTypeLikeUint_ = ...) -> NDArray[np.uint32 | np.uint64]: ...
  50. ###
  51. class ISeedSequence(abc.ABC):
  52. @abc.abstractmethod
  53. def generate_state(self, /, n_words: int, dtype: _DTypeLikeUint_ = ...) -> NDArray[np.uint32 | np.uint64]: ...
  54. class ISpawnableSeedSequence(ISeedSequence, abc.ABC):
  55. @abc.abstractmethod
  56. def spawn(self, /, n_children: int) -> list[Self]: ...
  57. class SeedlessSeedSequence(_GenerateStateMixin, ISpawnableSeedSequence):
  58. def spawn(self, /, n_children: int) -> list[Self]: ...
  59. class SeedSequence(_GenerateStateMixin, ISpawnableSeedSequence):
  60. __pyx_vtable__: ClassVar[CapsuleType] = ...
  61. entropy: int | Sequence[int] | None
  62. spawn_key: tuple[int, ...]
  63. pool_size: int
  64. n_children_spawned: int
  65. pool: NDArray[np.uint32]
  66. def __init__(
  67. self,
  68. /,
  69. entropy: _ArrayLikeInt_co | None = None,
  70. *,
  71. spawn_key: Sequence[int] = (),
  72. pool_size: int = 4,
  73. n_children_spawned: int = ...,
  74. ) -> None: ...
  75. def spawn(self, /, n_children: int) -> list[Self]: ...
  76. @property
  77. def state(self) -> _SeedSeqState: ...
  78. class BitGenerator(_CythonMixin, abc.ABC):
  79. lock: Lock
  80. @property
  81. def state(self) -> Mapping[str, Any]: ...
  82. @state.setter
  83. def state(self, value: Mapping[str, Any], /) -> None: ...
  84. @property
  85. def seed_seq(self) -> ISeedSequence: ...
  86. @property
  87. def ctypes(self) -> _Interface: ...
  88. @property
  89. def cffi(self) -> _Interface: ...
  90. @property
  91. def capsule(self) -> CapsuleType: ...
  92. #
  93. def __init__(self, /, seed: _ArrayLikeInt_co | SeedSequence | None = None) -> None: ...
  94. def __reduce__(self) -> tuple[Callable[[str], Self], tuple[str], tuple[Mapping[str, Any], ISeedSequence]]: ...
  95. def spawn(self, /, n_children: int) -> list[Self]: ...
  96. def _benchmark(self, /, cnt: int, method: str = "uint64") -> None: ...
  97. #
  98. @overload
  99. def random_raw(self, /, size: None = None, output: Literal[True] = True) -> int: ...
  100. @overload
  101. def random_raw(self, /, size: _ShapeLike, output: Literal[True] = True) -> NDArray[np.uint64]: ...
  102. @overload
  103. def random_raw(self, /, size: _ShapeLike | None, output: Literal[False]) -> None: ...
  104. @overload
  105. def random_raw(self, /, size: _ShapeLike | None = None, *, output: Literal[False]) -> None: ...