feat: default to PCG64 rng if numpy is installed (#7)

d20/__init__.py

@@ -1,5 +1,5 @@
 from . import diceast as ast
-from . import utils
+from . import rand, utils
 from .dice import *
 from .errors import *
 from .expression import *

d20/dice.py

@@ -1,10 +1,11 @@
+import random
 from enum import IntEnum
 from typing import Callable, Mapping, MutableMapping, Optional, Type, TypeVar, Union
 
 import cachetools
 import lark
 
-from . import diceast as ast, utils
+from . import diceast as ast, rand, utils
 from .errors import *
 from .expression import *
 from .stringifiers import MarkdownStringifier, Stringifier
@@ -133,7 +134,7 @@ def __repr__(self):
 class Roller:
     """The main class responsible for parsing dice into an AST and evaluating that AST."""
 
-    def __init__(self, context: Optional[RollContext] = None):
+    def __init__(self, context: RollContext = None, rng: random.Random = rand.random_impl):
         if context is None:
             context = RollContext()
 
@@ -150,7 +151,8 @@ def __init__(self, context: Optional[RollContext] = None):
             ast.Dice: self._eval_dice
         }
         self._parse_cache: MutableMapping[str, ASTNode] = cachetools.LFUCache(256)
-        self.context: RollContext = context
+        self.context = context
+        self.rng = rng
 
     def roll(self,
              expr: Union[str, ASTNode],
@@ -273,4 +275,4 @@ def _eval_operateddice(self, node: ast.OperatedDice) -> ExpressionNode:
         return self._eval_operatedset(node)
 
     def _eval_dice(self, node: ast.Dice) -> Dice:
-        return Dice.new(node.num, node.size, context=self.context)
+        return Dice.new(node.num, node.size, context=self.context, rng=self.rng)

d20/expression.py

@@ -1,8 +1,7 @@
 import abc
 import random
 
-from . import diceast as ast
-from . import errors
+from . import diceast as ast, errors, rand
 
 __all__ = (
     "Number", "Expression", "Literal", "UnOp", "BinOp", "Parenthetical", "Set", "Dice", "Die",
@@ -329,27 +328,29 @@ def __copy__(self):
 
 class Dice(Set):
     """A set of Die."""
-    __slots__ = ("num", "size", "_context")
+    __slots__ = ("num", "size", "_context", "_rng")
 
-    def __init__(self, num, size, values, operations=None, context=None, **kwargs):
+    def __init__(self, num, size, values, operations=None, context=None, rng=rand.random_impl, **kwargs):
         """
         :type num: int
         :type size: int|str
         :type values: list of Die
         :type operations: list[SetOperator]
         :type context: dice.RollContext
+        :type rng: random.Random
         """
         super().__init__(values, operations, **kwargs)
         self.num = num
         self.size = size
         self._context = context
+        self._rng = rng
 
     @classmethod
-    def new(cls, num, size, context=None):
-        return cls(num, size, [Die.new(size, context=context) for _ in range(num)], context=context)
+    def new(cls, num, size, context=None, rng=rand.random_impl):
+        return cls(num, size, [Die.new(size, context=context, rng=rng) for _ in range(num)], context=context, rng=rng)
 
     def roll_another(self):
-        self.values.append(Die.new(self.size, context=self._context))
+        self.values.append(Die.new(self.size, context=self._context, rng=self._rng))
 
     @property
     def children(self):
@@ -359,28 +360,30 @@ def __repr__(self):
         return f"<Dice num={self.num} size={self.size} values={self.values} operations={self.operations}>"
 
     def __copy__(self):
-        return Dice(num=self.num, size=self.size, context=self._context,
-                    values=self.values.copy(), operations=self.operations.copy(), )
+        return Dice(num=self.num, size=self.size, context=self._context, rng=self._rng,
+                    values=self.values.copy(), operations=self.operations.copy())
 
 
 class Die(Number):  # part of diceexpr
     """Represents a single die."""
-    __slots__ = ("size", "values", "_context")
+    __slots__ = ("size", "values", "_context", "_rng")
 
-    def __init__(self, size, values, context=None):
+    def __init__(self, size, values, context=None, rng=rand.random_impl):
         """
         :type size: int
         :type values: list of Literal
         :type context: dice.RollContext
+        :type rng: random.Random
         """
         super().__init__()
         self.size = size
         self.values = values
         self._context = context
+        self._rng = rng
 
     @classmethod
-    def new(cls, size, context=None):
-        inst = cls(size, [], context=context)
+    def new(cls, size, context=None, rng=rand.random_impl):
+        inst = cls(size, [], context=context, rng=rng)
         inst._add_roll()
         return inst
 
@@ -402,9 +405,9 @@ def _add_roll(self):
         if self._context:
             self._context.count_roll()
         if self.size == '%':
-            n = Literal(random.randrange(10) * 10)
+            n = Literal(self._rng.randrange(10) * 10)
         else:
-            n = Literal(random.randrange(self.size) + 1)  # 200ns faster than randint(1, self._size)
+            n = Literal(self._rng.randrange(self.size) + 1)  # 200ns faster than randint(1, self._size)
         self.values.append(n)
 
     def reroll(self):

d20/rand.py

@@ -0,0 +1,79 @@
+"""
+This module exposes a default implementation of the RNG based on the environment the library is installed in.
+
+In a thin installation, this will be MT2002 as implemented in the random stdlib module.
+If numpy is installed, this will be PCG64DXSM or PCG64 if available. If neither are available, it will fall back
+to whatever NumPy decides is a sane default for a bit generator
+(see https://numpy.org/doc/stable/reference/random/bit_generators/index.html).
+
+For more information, see https://github.com/avrae/d20/issues/7.
+
+Thanks to @posita for inspiring the implementation of NumpyRandom.
+"""
+
+import random
+from typing import Any, NewType, Optional, Sequence, TYPE_CHECKING, Union
+
+__all__ = ("random_impl",)
+
+_BitGenT = NewType('_BitGenT', Any)
+_SeedT = Optional[Union[int, Sequence[int]]]
+
+# the default random implementation - just use stdlib random (MT2002)
+random_impl = random.Random()
+
+# if np is installed and it has the random module, make use of PCG64
+# todo tests, docs
+try:
+    import numpy.random  # added in numpy 1.17
+    from numpy.random import Generator
+
+    if TYPE_CHECKING:
+        try:
+            # this was only exposed in numpy 1.19 - we only import these for type checking
+            # noinspection PyUnresolvedReferences
+            from numpy.random import BitGenerator, SeedSequence
+
+            _BitGenT = BitGenerator
+            _SeedT = Union[_SeedT, SeedSequence]
+        except ImportError:
+            pass
+
+
+    class NumpyRandom(random.Random):
+        def __init__(self, generator: _BitGenT, x: _SeedT = None):
+            self._gen = Generator(generator)
+            super().__init__(x)
+
+        def random(self) -> float:
+            return self._gen.random()
+
+        def getrandbits(self, k: int) -> int:
+            if k < 0:
+                raise ValueError('number of bits must be non-negative')
+            numbytes = (k + 7) // 8  # bits / 8 and rounded up
+            x = int.from_bytes(self._gen.bytes(numbytes), 'big')
+            return x >> (numbytes * 8 - k)  # trim excess bits
+
+        def seed(self, a: _SeedT = None, version: int = 2):
+            # note that this takes in a different type than random.Random.seed()
+            # this is because BitGenerator's seed requires an int/sequence of ints, while Random accepts
+            # floats, strs, etc; for the common case we expect the user to pass an int
+            bg_type = type(self._gen.bit_generator)
+            self._gen = Generator(bg_type(a))
+
+        def getstate(self):
+            return self._gen.bit_generator.state
+
+        def setstate(self, state):
+            self._gen.bit_generator.state = state
+
+
+    if hasattr(numpy.random, "PCG64DXSM"):  # available in numpy 1.21 and up
+        random_impl = NumpyRandom(numpy.random.PCG64DXSM())
+    elif hasattr(numpy.random, "PCG64"):  # available in numpy 1.17 and up
+        random_impl = NumpyRandom(numpy.random.PCG64())
+    elif hasattr(numpy.random, "default_rng"):
+        random_impl = NumpyRandom(numpy.random.default_rng().bit_generator)
+except ImportError:
+    pass

tests/conftest.py

@@ -1,10 +1,11 @@
-import random
+import d20
 
 import pytest
 
 
 @pytest.fixture(autouse=True, scope="function")
 def global_fixture():
     """Seed each individual test with the same seed, so that different runs of the same test are deterministic"""
-    random.seed(42)
+    # noinspection PyProtectedMember
+    d20._roller.rng.seed(42)
     yield