Hypothesis strategy for random units expressions

units_strategy.py

import hypothesis.strategies as s
from support.UnitRegistry import units, Q
from pint import UnitRegistry
from pint.errors import UndefinedUnitError
from pprint import PrettyPrinter
pp = PrettyPrinter(indent=2)


def _in_unit_registry(u: str, ureg: UnitRegistry) -> bool:
    try:
        _ = ureg.Quantity(42, u)
        return True
    except UndefinedUnitError:
        return False


_ok_units = list(filter(lambda u: _in_unit_registry(u, units), dir(units)))
# The following units are non-linear; we can't tolerate them. Degrees Kelvin
# are ok ('degK', 'kelvin')
_ok_units.remove('degC')
_ok_units.remove('degF')
_ok_units.remove('degR')
_ok_units.remove('celsius')
_ok_units.remove('rankine')
_ok_units.remove('fahrenheit')

RandomUnit = s.sampled_from(_ok_units)

# High powers will sometimes generate NaNs. Go ahead and enable them if your
# tests can tolerate NaNs from pint
RandomPowerOperator = s.sampled_from([
    "** -9", "** -8", "** -7", "** -6", "** -5", "** -4",
    "** -3", "** -2", "** -1", "** 2",  "** 3",  "** 4",
    "** 5", "** 6", "** 7", "** 8", "** 9", ])

# These lower powers seldom, if ever, produce NaNs.
RandomPowerOperator = s.sampled_from([
    "** -3", "** -2", "** -1", "** 2",  "** 3"])

RandomMultiplicativeOperator = s.sampled_from(["*", "/"])
RandomOperator = s.one_of(RandomPowerOperator,
                          RandomMultiplicativeOperator)
RandomUnitTupleExpression = s.recursive(
    RandomUnit, lambda children: s.one_of(
        children,
        s.tuples(children, RandomMultiplicativeOperator, children),
        s.tuples(children, RandomPowerOperator) ) )


def unit_tuple_expression_to_string(expr):
    if isinstance(expr, str):
        return expr
    elif isinstance(expr, tuple):
        children = [unit_tuple_expression_to_string(e) for e in expr]
        return "(" + " ".join(children) + ")"
    else:
        raise ValueError(f"""expr {expr} must be a string or a tuple, not a 
        {type(expr)}""")

# Uncomment for debugging
pp.pprint([Q(s.floats().example(),
             unit_tuple_expression_to_string(
                 RandomUnitTupleExpression.example()))
           for _ in range(20)])