1import random
2
3import numpy as np
4
5from numba.core import types
6from .templates import (ConcreteTemplate, AbstractTemplate, AttributeTemplate,
7                        CallableTemplate, Registry, signature)
8from numba.np.numpy_support import numpy_version
9
10
11registry = Registry()
12infer = registry.register
13infer_global = registry.register_global
14infer_getattr = registry.register_attr
15
16
17# random.random(), random.seed() etc. are not plain functions, they are bound
18# methods of a private object.  We have to be careful to use a well-known
19# object (e.g. the string "random.seed") as a key, not the bound method itself.
20# (same for np.random.random(), etc.)
21
22_int_types = sorted(set((types.intp, types.int64)))
23# Should we support float32?
24_float_types = [types.float64]
25
26
27#
28# Basics
29#
30
31def normalize_shape(shape):
32    if isinstance(shape, types.Integer):
33        return types.intp, 1
34    elif (isinstance(shape, types.BaseTuple) and
35          all(isinstance(v, types.Integer)) for v in shape):
36        ndim = len(shape)
37        return types.UniTuple(types.intp, ndim), ndim
38    else:
39        raise TypeError("invalid size type %s" % (shape,))
40
41
42class RandomTemplate(CallableTemplate):
43    """
44    A template helper to transparently handle the typing of array-returning
45    np.random.* functions.
46    """
47
48    def array_typer(self, scalar_typer, size=None):
49        prefix = self.key.split('.')[0]
50        assert prefix in ('np', 'random'), self.key
51
52        if size is None:
53            # Scalar variant
54            def typer(*args, **kwargs):
55                return scalar_typer(*args, **kwargs)
56        else:
57            # Array variant (only for the 'np.random.*' namespace)
58            def typer(*args, **kwargs):
59                if prefix == 'random':
60                    raise TypeError("unexpected size parameter for %r"
61                                    % (self.key,))
62                shape, ndim = normalize_shape(size)
63                # Type the scalar variant and wrap the result in an array
64                # of the appropriate dimensionality.
65                sig = scalar_typer(*args, **kwargs)
66                if sig is not None:
67                    return signature(
68                        types.Array(sig.return_type, ndim, 'C'),
69                        *(sig.args + (shape,)))
70
71        return typer
72
73
74class ConcreteRandomTemplate(RandomTemplate):
75    """
76    A RandomTemplate subclass using the `cases` attribute as a list of
77    allowed scalar signatures.
78    """
79
80    def array_typer(self, size=None):
81        key = self.key
82        cases = self.cases
83        context = self.context
84
85        def concrete_scalar_typer(*args, **kwargs):
86            # Filter out omitted args
87            while args and args[-1] is None:
88                args = args[:-1]
89            return context.resolve_overload(key, cases, args, kwargs)
90
91        return RandomTemplate.array_typer(self, concrete_scalar_typer, size)
92
93
94@infer_global(random.getrandbits, typing_key="random.getrandbits")
95class Random_getrandbits(ConcreteTemplate):
96    cases = [signature(types.uint64, types.int32)]
97
98@infer_global(random.random, typing_key="random.random")
99@infer_global(np.random.random, typing_key="np.random.random")
100class Random_random(ConcreteRandomTemplate):
101    cases = [signature(types.float64)]
102
103    def generic(self):
104        def typer(size=None):
105            return self.array_typer(size)()
106        return typer
107
108if numpy_version >= (1, 17):
109    infer_global(
110        np.random.random_sample,
111        typing_key="np.random.random_sample",
112    )(Random_random)
113    infer_global(
114        np.random.sample,
115        typing_key="np.random.sample",
116    )(Random_random)
117    infer_global(
118        np.random.ranf,
119        typing_key="np.random.ranf",
120    )(Random_random)
121
122
123@infer_global(random.randint, typing_key="random.randint")
124class Random_randint(ConcreteTemplate):
125    cases = [signature(tp, tp, tp) for tp in _int_types]
126
127@infer_global(np.random.randint, typing_key="np.random.randint")
128class Random_randint(ConcreteRandomTemplate):
129    cases = [signature(tp, tp) for tp in _int_types]
130    cases += [signature(tp, tp, tp) for tp in _int_types]
131
132    def generic(self):
133        def typer(low, high=None, size=None):
134            return self.array_typer(size)(low, high)
135        return typer
136
137
138@infer_global(random.randrange, typing_key="random.randrange")
139class Random_randrange(ConcreteTemplate):
140    cases = [signature(tp, tp) for tp in _int_types]
141    cases += [signature(tp, tp, tp) for tp in _int_types]
142    cases += [signature(tp, tp, tp, tp) for tp in _int_types]
143
144@infer_global(random.seed, typing_key="random.seed")
145@infer_global(np.random.seed, typing_key="np.random.seed")
146class Random_seed(ConcreteTemplate):
147    cases = [signature(types.void, types.uint32)]
148
149
150#
151# Distributions
152#
153
154@infer_global(np.random.geometric, typing_key="np.random.geometric")
155@infer_global(np.random.logseries, typing_key="np.random.logseries")
156@infer_global(np.random.zipf, typing_key="np.random.zipf")
157class Numpy_geometric(ConcreteRandomTemplate):
158    cases = [signature(types.int64, tp) for tp in _float_types]
159
160    def generic(self):
161        def typer(a, size=None):
162            return self.array_typer(size)(a)
163        return typer
164
165@infer_global(np.random.binomial, typing_key="np.random.binomial")
166@infer_global(np.random.negative_binomial,
167                          typing_key="np.random.negative_binomial")
168class Numpy_negative_binomial(ConcreteRandomTemplate):
169    cases = [signature(types.int64, types.int64, tp) for tp in _float_types]
170
171    def generic(self):
172        def typer(n, p, size=None):
173            return self.array_typer(size)(n, p)
174        return typer
175
176@infer_global(np.random.poisson, typing_key="np.random.poisson")
177class Numpy_poisson(ConcreteRandomTemplate):
178    cases = [signature(types.int64, tp) for tp in _float_types]
179    cases += [signature(types.int64)]
180
181    def generic(self):
182        def typer(lam=None, size=None):
183            return self.array_typer(size)(lam)
184        return typer
185
186@infer_global(np.random.exponential, typing_key="np.random.exponential")
187@infer_global(np.random.rayleigh, typing_key="np.random.rayleigh")
188class Numpy_exponential(ConcreteRandomTemplate):
189    cases = [signature(tp, tp) for tp in _float_types]
190    cases += [signature(tp) for tp in _float_types]
191
192    def generic(self):
193        def typer(scale=None, size=None):
194            return self.array_typer(size)(scale)
195        return typer
196
197@infer_global(np.random.hypergeometric, typing_key="np.random.hypergeometric")
198class Numpy_hypergeometric(ConcreteRandomTemplate):
199    cases = [signature(tp, tp, tp, tp) for tp in _int_types]
200
201    def generic(self):
202        def typer(ngood, nbad, nsample, size=None):
203            return self.array_typer(size)(ngood, nbad, nsample)
204        return typer
205
206@infer_global(np.random.laplace, typing_key="np.random.laplace")
207@infer_global(np.random.logistic, typing_key="np.random.logistic")
208@infer_global(np.random.lognormal, typing_key="np.random.lognormal")
209@infer_global(np.random.normal, typing_key="np.random.normal")
210class Numpy_normal(ConcreteRandomTemplate):
211    cases = [signature(tp, tp, tp) for tp in _float_types]
212    cases += [signature(tp, tp) for tp in _float_types]
213    cases += [signature(tp) for tp in _float_types]
214
215    def generic(self):
216        def typer(loc=None, scale=None, size=None):
217            return self.array_typer(size)(loc, scale)
218        return typer
219
220@infer_global(np.random.gamma, typing_key="np.random.gamma")
221class Numpy_gamma(ConcreteRandomTemplate):
222    cases = [signature(tp, tp, tp) for tp in _float_types]
223    cases += [signature(tp, tp) for tp in _float_types]
224
225    def generic(self):
226        def typer(shape, scale=None, size=None):
227            return self.array_typer(size)(shape, scale)
228        return typer
229
230@infer_global(np.random.triangular, typing_key="np.random.triangular")
231class Random_ternary_distribution(ConcreteRandomTemplate):
232    cases = [signature(tp, tp, tp, tp) for tp in _float_types]
233
234    def generic(self):
235        def typer(left, mode, right, size=None):
236            return self.array_typer(size)(left, mode, right)
237        return typer
238
239
240@infer_global(np.random.beta, typing_key="np.random.beta")
241@infer_global(np.random.f, typing_key="np.random.f")
242@infer_global(np.random.gumbel, typing_key="np.random.gumbel")
243@infer_global(np.random.uniform, typing_key="np.random.uniform")
244@infer_global(np.random.vonmises, typing_key="np.random.vonmises")
245@infer_global(np.random.wald, typing_key="np.random.wald")
246@infer_global(random.betavariate, typing_key="random.betavariate")
247@infer_global(random.gammavariate, typing_key="random.gammavariate")
248@infer_global(random.gauss, typing_key="random.gauss")
249@infer_global(random.lognormvariate, typing_key="random.lognormvariate")
250@infer_global(random.normalvariate, typing_key="random.normalvariate")
251@infer_global(random.uniform, typing_key="random.uniform")
252@infer_global(random.vonmisesvariate, typing_key="random.vonmisesvariate")
253@infer_global(random.weibullvariate, typing_key="random.weibullvariate")
254class Random_binary_distribution(ConcreteRandomTemplate):
255    cases = [signature(tp, tp, tp) for tp in _float_types]
256
257    def generic(self):
258        def typer(a, b, size=None):
259            return self.array_typer(size)(a, b)
260        return typer
261
262
263@infer_global(np.random.chisquare, typing_key="np.random.chisquare")
264@infer_global(np.random.pareto, typing_key="np.random.pareto")
265@infer_global(np.random.power, typing_key="np.random.power")
266@infer_global(np.random.standard_gamma, typing_key="np.random.standard_gamma")
267@infer_global(np.random.standard_t, typing_key="np.random.standard_t")
268@infer_global(np.random.weibull, typing_key="np.random.weibull")
269@infer_global(random.expovariate, typing_key="random.expovariate")
270@infer_global(random.paretovariate, typing_key="random.paretovariate")
271class Random_unary_distribution(ConcreteRandomTemplate):
272    cases = [signature(tp, tp) for tp in _float_types]
273
274    def generic(self):
275        def typer(a, size=None):
276            return self.array_typer(size)(a)
277        return typer
278
279
280@infer_global(np.random.standard_cauchy,
281                          typing_key="np.random.standard_cauchy")
282@infer_global(np.random.standard_normal,
283                          typing_key="np.random.standard_normal")
284@infer_global(np.random.standard_exponential,
285                          typing_key="np.random.standard_exponential")
286class Random_nullary_distribution(ConcreteRandomTemplate):
287    cases = [signature(tp) for tp in _float_types]
288
289    def generic(self):
290        def typer(size=None):
291            return self.array_typer(size)()
292        return typer
293
294
295@infer_global(random.triangular, typing_key="random.triangular")
296class Random_triangular(ConcreteTemplate):
297    cases = [signature(tp, tp, tp) for tp in _float_types]
298    cases += [signature(tp, tp, tp, tp) for tp in _float_types]
299
300# NOTE: some functions can have @overloads in numba.targets.randomimpl,
301# and therefore don't need a typing declaration here.
302