Skip to content

Commit a53af6c

Browse files
dpgeorgedpgeorge
committed
extmod/modurandom: Add some extra random functions.
Functions added are: - randint - randrange - choice - random - uniform They are enabled with configuration variable MICROPY_PY_URANDOM_EXTRA_FUNCS, which is disabled by default. It is enabled for unix coverage build and stmhal.
1 parent f5c42dd commit a53af6c

File tree

5 files changed

+144
-1
lines changed

5 files changed

+144
-1
lines changed

‎extmod/modurandom.c‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,24 @@ STATIC uint32_t yasmarang(void)
5353

5454
// End of Yasmarang
5555

56+
#if MICROPY_PY_URANDOM_EXTRA_FUNCS
57+
58+
// returns an unsigned integer below the given argument
59+
// n must not be zero
60+
STATIC uint32_t yasmarang_randbelow(uint32_t n) {
61+
uint32_t mask = 1;
62+
while ((n & mask) < n) {
63+
mask = (mask << 1) | 1;
64+
}
65+
uint32_t r;
66+
do {
67+
r = yasmarang() & mask;
68+
} while (r >= n);
69+
return r;
70+
}
71+
72+
#endif
73+
5674
STATIC mp_obj_t mod_urandom_getrandbits(mp_obj_t num_in) {
5775
int n = mp_obj_get_int(num_in);
5876
if (n > 32 || n == 0) {
@@ -75,10 +93,122 @@ STATIC mp_obj_t mod_urandom_seed(mp_obj_t seed_in) {
7593
}
7694
STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_urandom_seed_obj, mod_urandom_seed);
7795

96+
#if MICROPY_PY_URANDOM_EXTRA_FUNCS
97+
98+
STATIC mp_obj_t mod_urandom_randrange(size_t n_args, const mp_obj_t *args) {
99+
mp_int_t start = mp_obj_get_int(args[0]);
100+
if (n_args == 1) {
101+
// range(stop)
102+
if (start > 0) {
103+
return mp_obj_new_int(yasmarang_randbelow(start));
104+
} else {
105+
nlr_raise(mp_obj_new_exception(&mp_type_ValueError));
106+
}
107+
} else {
108+
mp_int_t stop = mp_obj_get_int(args[1]);
109+
if (n_args == 2) {
110+
// range(start, stop)
111+
if (start < stop) {
112+
return mp_obj_new_int(start + yasmarang_randbelow(stop - start));
113+
} else {
114+
nlr_raise(mp_obj_new_exception(&mp_type_ValueError));
115+
}
116+
} else {
117+
// range(start, stop, step)
118+
mp_int_t step = mp_obj_get_int(args[2]);
119+
mp_int_t n;
120+
if (step > 0) {
121+
n = (stop - start + step - 1) / step;
122+
} else if (step < 0) {
123+
n = (stop - start + step + 1) / step;
124+
} else {
125+
nlr_raise(mp_obj_new_exception(&mp_type_ValueError));
126+
}
127+
if (n > 0) {
128+
return mp_obj_new_int(start + step * yasmarang_randbelow(n));
129+
} else {
130+
nlr_raise(mp_obj_new_exception(&mp_type_ValueError));
131+
}
132+
}
133+
}
134+
}
135+
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_urandom_randrange_obj, 1, 3, mod_urandom_randrange);
136+
137+
STATIC mp_obj_t mod_urandom_randint(mp_obj_t a_in, mp_obj_t b_in) {
138+
mp_int_t a = mp_obj_get_int(a_in);
139+
mp_int_t b = mp_obj_get_int(b_in);
140+
if (a <= b) {
141+
return mp_obj_new_int(a + yasmarang_randbelow(b - a + 1));
142+
} else {
143+
nlr_raise(mp_obj_new_exception(&mp_type_ValueError));
144+
}
145+
}
146+
STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_urandom_randint_obj, mod_urandom_randint);
147+
148+
STATIC mp_obj_t mod_urandom_choice(mp_obj_t seq) {
149+
mp_int_t len = mp_obj_get_int(mp_obj_len(seq));
150+
if (len > 0) {
151+
return mp_obj_subscr(seq, mp_obj_new_int(yasmarang_randbelow(len)), MP_OBJ_SENTINEL);
152+
} else {
153+
nlr_raise(mp_obj_new_exception(&mp_type_IndexError));
154+
}
155+
}
156+
STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_urandom_choice_obj, mod_urandom_choice);
157+
158+
#if MICROPY_PY_BUILTINS_FLOAT
159+
160+
// returns a number in the range [0..1) using Yasmarang to fill in the fraction bits
161+
STATIC mp_float_t yasmarang_float(void) {
162+
#if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_DOUBLE
163+
typedef uint64_t mp_float_int_t;
164+
#elif MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
165+
typedef uint32_t mp_float_int_t;
166+
#endif
167+
union {
168+
mp_float_t f;
169+
#if MP_ENDIANNESS_LITTLE
170+
struct { mp_float_int_t frc:MP_FLOAT_FRAC_BITS, exp:MP_FLOAT_EXP_BITS, sgn:1; } p;
171+
#else
172+
struct { mp_float_int_t sgn:1, exp:MP_FLOAT_EXP_BITS, frc:MP_FLOAT_FRAC_BITS; } p;
173+
#endif
174+
} u;
175+
u.p.sgn = 0;
176+
u.p.exp = (1 << (MP_FLOAT_EXP_BITS - 1)) - 1;
177+
if (MP_FLOAT_FRAC_BITS <= 32) {
178+
u.p.frc = yasmarang();
179+
} else {
180+
u.p.frc = ((uint64_t)yasmarang() << 32) | (uint64_t)yasmarang();
181+
}
182+
return u.f - 1;
183+
}
184+
185+
STATIC mp_obj_t mod_urandom_random(void) {
186+
return mp_obj_new_float(yasmarang_float());
187+
}
188+
STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_urandom_random_obj, mod_urandom_random);
189+
190+
STATIC mp_obj_t mod_urandom_uniform(mp_obj_t a_in, mp_obj_t b_in) {
191+
mp_float_t a = mp_obj_get_float(a_in);
192+
mp_float_t b = mp_obj_get_float(b_in);
193+
return mp_obj_new_float(a + (b - a) * yasmarang_float());
194+
}
195+
STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_urandom_uniform_obj, mod_urandom_uniform);
196+
197+
#endif
198+
199+
#endif // MICROPY_PY_URANDOM_EXTRA_FUNCS
200+
78201
STATIC const mp_rom_map_elem_t mp_module_urandom_globals_table[] = {
79202
{ MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_urandom) },
80203
{ MP_ROM_QSTR(MP_QSTR_getrandbits), MP_ROM_PTR(&mod_urandom_getrandbits_obj) },
81204
{ MP_ROM_QSTR(MP_QSTR_seed), MP_ROM_PTR(&mod_urandom_seed_obj) },
205+
#if MICROPY_PY_URANDOM_EXTRA_FUNCS
206+
{ MP_ROM_QSTR(MP_QSTR_randrange), MP_ROM_PTR(&mod_urandom_randrange_obj) },
207+
{ MP_ROM_QSTR(MP_QSTR_randint), MP_ROM_PTR(&mod_urandom_randint_obj) },
208+
{ MP_ROM_QSTR(MP_QSTR_choice), MP_ROM_PTR(&mod_urandom_choice_obj) },
209+
{ MP_ROM_QSTR(MP_QSTR_random), MP_ROM_PTR(&mod_urandom_random_obj) },
210+
{ MP_ROM_QSTR(MP_QSTR_uniform), MP_ROM_PTR(&mod_urandom_uniform_obj) },
211+
#endif
82212
};
83213

84214
STATIC MP_DEFINE_CONST_DICT(mp_module_urandom_globals, mp_module_urandom_globals_table);

‎py/mpconfig.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -787,6 +787,11 @@ typedef double mp_float_t;
787787
#define MICROPY_PY_URANDOM (0)
788788
#endif
789789

790+
// Whether to include: randrange, randint, choice, random, uniform
791+
#ifndef MICROPY_PY_URANDOM_EXTRA_FUNCS
792+
#define MICROPY_PY_URANDOM_EXTRA_FUNCS (0)
793+
#endif
794+
790795
#ifndef MICROPY_PY_MACHINE
791796
#define MICROPY_PY_MACHINE (0)
792797
#endif

‎py/qstrdefs.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -686,4 +686,11 @@ Q(dupterm)
686686
Q(urandom)
687687
Q(getrandbits)
688688
Q(seed)
689+
#if MICROPY_PY_URANDOM_EXTRA_FUNCS
690+
Q(randrange)
691+
Q(randint)
692+
Q(choice)
693+
Q(random)
694+
Q(uniform)
695+
#endif
689696
#endif

‎stmhal/mpconfigport.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
#define MICROPY_PY_IO_FILEIO (1)
8181
#define MICROPY_PY_UBINASCII (1)
8282
#define MICROPY_PY_URANDOM (1)
83+
#define MICROPY_PY_URANDOM_EXTRA_FUNCS (1)
8384
#define MICROPY_PY_UCTYPES (1)
8485
#define MICROPY_PY_UZLIB (1)
8586
#define MICROPY_PY_UJSON (1)

‎unix/Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ freedos:
193193

194194
# build an interpreter for coverage testing and do the testing
195195
coverage:
196-
$(MAKE) COPT="-O0" CFLAGS_EXTRA='-fprofile-arcs -ftest-coverage -Wdouble-promotion -Wformat -Wmissing-declarations -Wmissing-prototypes -Wold-style-definition -Wpointer-arith -Wshadow -Wsign-compare -Wuninitialized -Wunused-parameter -DMICROPY_UNIX_COVERAGE' LDFLAGS_EXTRA='-fprofile-arcs -ftest-coverage' BUILD=build-coverage PROG=micropython_coverage
196+
$(MAKE) COPT="-O0" CFLAGS_EXTRA='-fprofile-arcs -ftest-coverage -Wdouble-promotion -Wformat -Wmissing-declarations -Wmissing-prototypes -Wold-style-definition -Wpointer-arith -Wshadow -Wsign-compare -Wuninitialized -Wunused-parameter -DMICROPY_UNIX_COVERAGE -DMICROPY_PY_URANDOM_EXTRA_FUNCS' LDFLAGS_EXTRA='-fprofile-arcs -ftest-coverage' BUILD=build-coverage PROG=micropython_coverage
197197

198198
coverage_test: coverage
199199
$(eval DIRNAME=$(notdir $(CURDIR)))

0 commit comments

Comments
 (0)