codekingpro/portable-devtools

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