diff --git a/extmod/modurandom.c b/extmod/modurandom.c index f44510be9b..21fbac6942 100644 --- a/extmod/modurandom.c +++ b/extmod/modurandom.c @@ -87,9 +87,12 @@ STATIC uint32_t yasmarang_randbelow(uint32_t n) { STATIC mp_obj_t mod_urandom_getrandbits(mp_obj_t num_in) { int n = mp_obj_get_int(num_in); - if (n > 32 || n == 0) { + if (n > 32 || n < 0) { mp_raise_ValueError(MP_ERROR_TEXT("bits must be 32 or less")); } + if (n == 0) { + return MP_OBJ_NEW_SMALL_INT(0); + } uint32_t mask = ~0; // Beware of C undefined behavior when shifting by >= than bit size mask >>= (32 - n); diff --git a/tests/extmod/urandom_basic.py b/tests/extmod/urandom_basic.py index 180197398f..f7f5a6d691 100644 --- a/tests/extmod/urandom_basic.py +++ b/tests/extmod/urandom_basic.py @@ -22,8 +22,11 @@ r = random.getrandbits(16) random.seed(1) print(random.getrandbits(16) == r) -# check that it throws an error for zero bits +# check that zero bits works +print(random.getrandbits(0)) + +# check that it throws an error for negative bits try: - random.getrandbits(0) + random.getrandbits(-1) except ValueError: print("ValueError") diff --git a/tests/extmod/urandom_basic.py.exp b/tests/extmod/urandom_basic.py.exp new file mode 100644 index 0000000000..d629828d79 --- /dev/null +++ b/tests/extmod/urandom_basic.py.exp @@ -0,0 +1,4 @@ +True +True +0 +ValueError