import random import collections import itertools import re basestr = (str,unicode) class DieJar(object): parser = re.compile(r'^(?:_?(?P<num>[1-9][0-9]*))?d(?P<sides>[1-9][0-9]*)$') def __getattr__(self, key): out = [] for sect in key.split('__'): out.extend(self.parse_section(sect)) return Dice(out) def parse_section(self, key): out = self.parser.match(key) if out is None: raise AttributeError('invalid die specification') else: dct = out.groupdict() dct['sides'] = int(dct['sides']) dct['num'] = int(dct['num'] or 1) dice = [Die(dct['sides']) for __ in range(dct['num'])] return dice def iterable(obj): return (not isinstance(obj, basestr)) and isinstance(obj, collections.Iterable) def flatten(lis): return itertools.chain(*[(x if iterable(x) else [x]) for x in lis]) class Die(object): 'Note: hashed by sides, min and step. Consequently not necessarily preserved when used as a dictionary key' def __init__(self, sides, min=1, step=1): self.sides = sides self.min = min self.step = 1 self.sides = sides self.choices = range(min, (min+sides)*step, step) self._value = None self.combine_func = lambda a,b: a+b def roll(self): self._value = random.choice(self.choices) return self._value @property def value(self): if self._value is None: self.roll() return self._value def combine(self, other): if hasattr(other, 'value'): other = other.value return self.combine_func(self.value, other.value) def __str__(self): base = 'd%d' % self.sides if self.min != 1: base = '%s+%d' % (base,self.min) if self.step != 1: base = '(%s)*%d' % (base, self.step) return base def __eq__(self, other): return (self.sides == other.sides) and (self.min == other.min) and (self.step == other.step) def __hash__(self): return hash((self.sides,self.min,self.step)) class Dice(collections.Sequence): 'A collection of dice, can be initialized either with Die instances or lists of Die instances' def __init__(self, *dice, **kw): self.dice = list(flatten(dice)) self.combiner = kw.get('combiner', lambda a,b:a+b) def __getitem__(self, k): return self.dice[k] def __len__(self): return len(self.dice) def roll(self): return reduce(self.combiner, (die.roll() for die in self.dice)) def __str__(self): groups = collections.Counter(self.dice) out = [] dice = sorted(groups, key=lambda k:-groups[k]) if len(dice) > 1: for die in dice[:-1]: count = groups[die] out.append('%d%s' % (count,die)) out = ','.join(out) out = ' and '.join([out, str(dice[-1])]) else: out = '%d%s' % (groups[dice[0]],dice[0]) return out MULT=1 ADD=2 class DieRoll(object): def __init__(self, dice=None, adjustments=None): self.dice = dice self.adjustments = [] if adjustments is None else adjustments def add_adjustment(self, add=None, mult=None): if None not in {add,mult}: raise ValueError('can\'t add two adjustments at once') elif {add,mult} - {None} == set(): raise ValueError('No adjustment given') if add is not None: adjustment = (ADD,add) elif mult is not None: adjustment = (MULT,mult) self.adjustments.append(adjustment) def roll(self): result = self.dice.roll() for type,bonus in self.adjustments: if type == MULT: result *= type elif type == ADD: result += bonus return result if __name__ == '__main__': import unittest tests = unittest.TestSuite() inst = lambda a:a() class TestDie(unittest.TestCase): def test_roll_plain(self): a = Die(6) for __ in range(40): self.assertLess(a.roll(), 7) self.assertGreater(a.roll(), 0) def test_roll_min(self): a = Die(6,min=4) for __ in range(40): self.assertLess(a.roll(), 6+4+1) self.assertGreater(a.roll(), 3) def test_roll_step(self): a = Die(6,step=2) for __ in range(40): self.assertLess(a.roll(), 6+(6*2)+1) self.assertGreater(a.roll(), 0) self.assertTrue((a.roll()-1) % 2 == 0) def test_str(self): self.assertEqual(str(Die(6)), 'd6') self.assertEqual(str(Die(6,min=2)), 'd6+2') class TestDice(unittest.TestCase): def test_init(self): dice = [Die(6) for __ in range(20)] a = Dice(dice) self.assertEqual(dice,a.dice) a = Dice(*dice) self.assertEqual(dice,a.dice) def test_roll(self): dice = [Die(6) for __ in range(2)] dice = Dice(dice) self.assertGreater(dice.roll(), 1) self.assertLess(dice.roll(), 13) class TestDieRoll(unittest.TestCase): pass unittest.main()