git.fiddlerwoaroof.com
libs/dice.py
3d5a515d
 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()