Improved random.randstr (preimport)

docs/pages/enhancements.md

@@ -82,10 +82,10 @@ A context manager is also available:
 
 `itertools` is extended with the following items:
 
-- `product2`: this is an improvement of the original `product`, also handling generators
-- `reset`: given a generator function decorated by `resettable`, this functions can reset a generator instantiated by this function
-- `resettable`: decorator for registering the reference to the generator function and its arguments used to make a generator, then making resettable each generator made by this function
-- `NonResettableGeneratorException`: specific exception for handling a generator not decorated by `resettable` thrown while trying to reset it with the `reset` function
+- `product2`: this is an improvement of the original `product`, also handling generators.
+- `reset`: given a generator function decorated by `resettable`, this functions can reset a generator instantiated by this function.
+- `resettable`: decorator for registering the reference to the generator function and its arguments used to make a generator, then making resettable each generator made by this function.
+- `NonResettableGeneratorException`: specific exception for handling a generator not decorated by `resettable` thrown while trying to reset it with the `reset` function.
 
 -----
 
@@ -114,8 +114,8 @@ A context manager is also available:
 
 `random` is slightly enhanced with a few new items:
 
-- `choice`: redefined to add an argument for an exclusion list (aim is to provide a short form instead of using list comprehension) and an extra argument for setting if an error shall be thrown when the resulting list is empty
-- `randstr`: allows to generate a random string with a given length (8 by default) and alphabet
+- `choice`: redefined to add an argument for an exclusion list (aim is to provide a short form instead of using list comprehension) and an extra argument for setting if an error shall be thrown when the resulting list is empty.
+- `randstr`: allows to generate a random string with a given length (8 by default) and alphabet ; it also supports a 'balance' parameter that ensures that there is no character that can have more than n / (n_alphabet - 1) occurrences and a 'blocksize' parameter to enforce balancing on a per-block basis.
 - `LFSR`: adds an implementation of the Linear-Feedback Shifting Register stream generator, with the possibility of recovering its parameters by setting a target and using the Berlekamp-Massey algorithm.
 - `Geffe`: adds an implementation of the Geffe stream generator.
 
@@ -125,10 +125,10 @@ A context manager is also available:
 
 `re` is enhanced with some new (fully lazy) functions to generate strings from regular expression patterns:
 
-- `randstr`: generates a single random string from the input regex
-- `randstrs`: provides a generator of N random strings from the input regex
-- `size`: computes the number of all possible strings from the input regex
-- `strings`: generates all possible strings from the input regex
+- `randstr`: generates a single random string from the input regex.
+- `randstrs`: provides a generator of N random strings from the input regex.
+- `size`: computes the number of all possible strings from the input regex.
+- `strings`: generates all possible strings from the input regex.
 
 -----
 
@@ -137,8 +137,8 @@ A context manager is also available:
 `string` is slightly enhanced with a few new functions:
 
 - `shorten`: shortens a string, taking by default the terminal width, otherwise a length of 40 characters (unless user-defined), and using an end token (by default "`...`").
-- `sort_natural`: sort a list of strings taking numbers into account (returns nothing)
-- `sorted_natural`: return a list of strings taking numbers into account
+- `sort_natural`: sort a list of strings taking numbers into account (returns nothing).
+- `sorted_natural`: return a list of strings taking numbers into account.
 
 -----
 

src/tinyscript/VERSION.txt

@@ -1 +1 @@
-1.30.2
+1.30.4

src/tinyscript/preimports/rand.py

@@ -19,15 +19,39 @@ def __choice(lst, exclusions=(), error=True):
 random.choice = __choice
 
 
-def __randstr(n=8, alphabet=string.ascii_lowercase+string.ascii_uppercase+string.digits):
-    """ Compose a random string of the given length with the given alphabet. """
+def __randstr(n=8, alphabet=string.ascii_lowercase+string.ascii_uppercase+string.digits, balance=False, blocksize=0):
+    """ Compose a random string of the given length with the given alphabet. It can be chosen if it has to be balanced,
+         either in its whole or per block (given a block size). Note that, when balancing per block, it is not ensured
+         that the whole string is balanced too. """
+    na = len(alphabet)
     if n < 0:
         raise ValueError("Bad random string length")
-    if len(alphabet) == 0:
+    if na == 0:
         raise ValueError("Bad alphabet")
-    s = ""
+    is_b = isinstance(alphabet, bytes)
+    s = ["", b""][is_b]
+    if is_b:
+        alphabet = [alphabet[i:i+1] for i in range(na)]
+    if balance:
+        bs = min(n, blocksize) or n
+        t = bs / (na-1 or 1)
+        if bs <= (na-1)/(1-(na-1)/na):
+            t = bs / (na-2)
+        orig_alphabet = alphabet[:] if is_b else alphabet
     for i in range(n):
-        s += random.choice(alphabet)
+        if balance:
+            if i == 0 or blocksize > 0 and i % blocksize == 0:
+                alphabet, cnts = orig_alphabet[:] if is_b else orig_alphabet, {}
+            while cnts.get(c := random.choice(alphabet), 0) >= t - 1:
+                if is_b:
+                    alphabet.remove(c)
+                else:
+                    alphabet = alphabet.replace(c, "")
+            cnts.setdefault(c, 0)
+            cnts[c] += 1
+        else:
+            c = random.choice(alphabet)
+        s += c
     return s
 random.randstr = __randstr
 

tests/test_preimports_random.py

@@ -2,6 +2,8 @@
 """Preimports random assets' tests.
 
 """
+from collections import Counter
+from math import log
 from tinyscript.preimports import random
 
 from utils import *
@@ -19,6 +21,12 @@ def test_utility_functions(self):
         self.assertNotIn("e", random.randstr(alphabet="abcd"))
         self.assertRaises(ValueError, random.randstr, -1)
         self.assertRaises(ValueError, random.randstr, 8, "")
+        self.assertTrue(isinstance(random.randstr(16, b"\x00\x01\x02"), bytes))
+        for n, na in zip([8, 16, 64], [3, 4, 5]):
+            for bs in [0, 16, 256]:
+                for i in range(512):
+                    self.assertLess(max(Counter(random.randstr(n, "".join(chr(i) for i in range(na)), True, bs))\
+                                        .values()), n/(na-1) if (bs or n) > (na-1)/(1-(na-1)/na) else n/(na-2))
 
     def test_random_lfsr(self):
         l = random.LFSR(target="0123456789abcdef")