import string import numpy as np from hypothesis import strategies as st from hypothesis.extra.numpy import arrays, array_shapes def _strategy_2d_array(dtype, minval=0, maxval=None, **kwargs): if 'min_side' in kwargs: min_side = kwargs.pop('min_side') else: min_side = 1 if 'max_side' in kwargs: max_side = kwargs.pop('max_side') else: max_side = None if dtype is np.int: elems = st.integers(minval, maxval, **kwargs) elif dtype is np.float: elems = st.floats(minval, maxval, **kwargs) elif dtype is np.str: elems = st.text(min_size=minval, max_size=maxval, **kwargs) else: raise ValueError('no elements strategy for dtype', dtype) return arrays(dtype, array_shapes(2, 2, min_side, max_side), elements=elems) def strategy_dtm(): return _strategy_2d_array(np.int, 0, 10000) def strategy_dtm_small(): return _strategy_2d_array(np.int, 0, 10, min_side=2, max_side=10) def strategy_2d_prob_distribution(): return _strategy_2d_array(np.float, 0, 1, allow_nan=False, allow_infinity=False) def strategy_tokens(*args, **kwargs): return st.lists(st.lists(st.text(*args, **kwargs))) def strategy_texts(*args, **kwargs): return st.lists(st.text(*args, **kwargs)) def strategy_texts_printable(*args, **kwargs): return strategy_texts(string.printable)