Add docstrings and doctests to deck.py + Python cleanups.
[poker.git] / deck.py
diff --git a/deck.py b/deck.py
new file mode 100644 (file)
index 0000000..4f1ca92
--- /dev/null
+++ b/deck.py
@@ -0,0 +1,224 @@
+"""Define a deck of cards, single-player scoring rules, and pretty-printing.
+"""
+
+from combinations import xunique_combinations
+
+
+SUITS = ['spades', 'hearts', 'diamonds', 'clubs']
+"""Ordered list of suits.
+"""
+
+FACE_CARDS = ['ten', 'jack', 'queen', 'king', 'ace']
+"""Ordered list of face cards.
+"""
+
+DECK = dict([(c/4, c%4) for c in range(52)])
+"""Cards in the deck are stored as (n,m), where `n` is the rank number
+`n=[0,12]`, and `m` is the suit `m=[0,3]`.
+
+Examples: 
+ (0,1)  is a two of hearts, 
+ (11,0) is a king of spades,
+ (12,3) is an ace of clubs
+"""
+
+
+def pp_card(card):
+    """Pretty-print a card.
+
+    >>> pp_card((0, 1))
+    '2h'
+    >>> pp_card((8, 3))
+    'Tc'
+    >>> pp_card((11, 0))
+    'Ks'
+    """
+    s = []
+    rank_num = card[0] + 2
+    if rank_num >= 10:
+        rank = FACE_CARDS[rank_num-10].capitalize()[0]
+    else:
+        rank = str(rank_num)
+    
+    suit = SUITS[card[1]]
+    return '%s%s' % (rank, suit[0])
+
+def unpp_card(card):
+    """Un-pretty-print a card.
+
+    >>> unpp_card('2h')
+    (0, 1)
+    >>> unpp_card('Tc')
+    (8, 3)
+    >>> unpp_card('Ks')
+    (11, 0)
+    """
+    rank,suit = card
+    try:
+        rank = int(rank)
+    except ValueError:
+        rank = rank.lower()
+        rank = [i for i,r in enumerate(FACE_CARDS) if r.startswith(rank)][0]
+        rank += 10
+    rank -= 2
+    suit = [i for i,s in enumerate(SUITS) if s[0] == suit][0]
+    return (rank, suit)
+
+def pp_hand(hand):
+    """Return a hand in a human readable format (pretty-print)
+
+    >>> pp_hand([(0, 1), (11, 0), (12, 3)])
+    '2h Ks Ac'
+    """
+    return ' '.join([pp_card(c) for c in hand])
+
+
+class FiveCardHand (object):
+    """Poker hand with five cards.
+
+    To determine which five card hand is stronger according to
+    standard poker rules, each hand is given a score and a pair
+    rank. The score cooresponds to the poker hand (Straight Flush=8,
+    Four of a Kind=7, etc...). The pair score is used to break ties.
+    The pair score creates a set of pairs then rank sorts within each
+    set.
+
+    >>> h = FiveCardHand([unpp_card(c) for c in ['7h','5h','7s','5s','Ks']])
+    >>> h.score()
+    (2, [(2, [3, 5]), (1, [11])])
+    >>> h = FiveCardHand([unpp_card(c) for c in ['8h','8d','5s','As','5c']])
+    >>> h.score()
+    (2, [(2, [3, 6]), (1, [12])])
+    >>> h = FiveCardHand([unpp_card(c) for c in ['Ah','Kd','Qs','Ks','2c']])
+    >>> h.score()
+    (1, [(2, [11]), (1, [0, 10, 12])])
+    >>> h = FiveCardHand([unpp_card(c) for c in ['Ah','5d','4s','3s','2c']])
+    >>> h.score()
+    (4, [(1, [0, 1, 2, 3, 12])])
+    >>> h.pp_score()
+    'straight - Ah 5d 4s 3s 2c'
+    """
+    types = ['x high', 'pair 2', 'double pair 2', 'pair 3', 'straight',
+             'flush', 'full house', 'pair 4', 'straight flush']
+
+    def __init__(self, hand=None):
+        self.hand = hand
+
+    def __str__(self):
+        return pp_hand(self.hand)
+
+    def __cmp__(self, other):
+        return cmp(self.score(), other.score())
+
+    def _pair_score(self, ranks):
+        """Returns the pair scores from a list of card ranks in a hand.
+
+        For the sake of clarity, name the pairscore tuple elements
+        `(paircount,ranks)`.
+
+        >>> h = FiveCardHand()
+        >>> h._pair_score([5, 3, 5, 3, 11])
+        [(2, [3, 5]), (1, [11])]
+        >>> h._pair_score([8, 8, 5, 8, 5])
+        [(3, [8]), (2, [5])]
+        """
+        rank_counts = [(ranks.count(r), r) for r in set(ranks)]
+        counts = sorted(set([rc[0] for rc in rank_counts]), reverse=True)
+        return [(c, [rc[1] for rc in rank_counts if rc[0] == c])
+                for c in counts]
+
+    def _is_full_house(self, pair_score, **kwargs):
+        """
+        >>> FiveCardHand()._is_full_house([(3, [8]), (2, [5])])
+        True
+        """
+        return [p[0] for p in pair_score] == [3, 2]
+
+    def _is_pair_4(self, pair_score, **kwargs):
+        return [p[0] for p in pair_score] == [4, 1]
+
+    def _is_pair_3(self, pair_score, **kwargs):
+        return [p[0] for p in pair_score] == [3, 1]
+
+    def _is_double_pair_2(self, pair_score, **kwargs):
+        return ([p[0] for p in pair_score] == [2, 1]
+                and len(pair_score[0][1]) == 2)
+
+    def _is_pair_2(self, pair_score, **kwargs):
+        return ([p[0] for p in pair_score] == [2, 1]
+                and len(pair_score[0][1]) == 1)
+
+    def _is_flush(self, suits, **kwargs):
+        return len(set(suits)) == 1
+
+    def _is_straight(self, ranks, **kwargs):  # Handles the low Ace as well
+        ranks = sorted(ranks)
+        diff  = [(x2-x1) for x1,x2 in zip(ranks, ranks[1:])]
+        return set(diff) == set([1]) or ranks == [0,1,2,3,12]
+
+    def _is_straight_flush(self, **kwargs):
+        return self._is_straight(**kwargs) and self._is_flush(**kwargs)
+
+    def _type_index(self, ranks, suits, pair_score):
+        for i,type in enumerate(reversed(self.types)):
+            if type == 'x high': continue
+            is_type = getattr(self, '_is_%s' % type.replace(' ', '_'))
+            if is_type(ranks=ranks, suits=suits, pair_score=pair_score):
+                return len(self.types) - i - 1
+        return 0
+
+    def score(self):
+        if hasattr(self, '_score'): return self._score # return cached
+        ranks,suits = zip(*self.hand)
+        pair_score = self._pair_score(ranks)
+        type_index = self._type_index(ranks, suits, pair_score)
+        self._score = (type_index, pair_score) # cache
+        return self._score
+
+    def pp_score(self):
+        score = self.score()
+        type = self.types[score[0]]
+        return '%s - %s' % (type, str(self))
+
+
+class SevenChooseFiveHand (FiveCardHand):
+    """Poker hand with seven cards.
+
+    To determine who wins a poker hand each player determines the best
+    hand they can make within their (7 choose 5) = 21 possible five card
+    combinations.
+
+    >>> h = SevenChooseFiveHand([unpp_card(c) for c in
+    ...     ['7h','5h','7s','5s','Ks', '7d', 'Jh']])
+    >>> h.full_hand
+    [(5, 1), (3, 1), (5, 0), (3, 0), (11, 0), (5, 2), (9, 1)]
+    >>> h.hand
+    [(5, 1), (3, 1), (5, 0), (3, 0), (5, 2)]
+    >>> h.score()
+    (6, [(3, [5]), (2, [3])])
+    >>> h.pp_score()
+    'full house - 7h 5h 7s 5s 7d Ks Jh'
+    """
+    # Prevent a bajillion calls to xunqiue_combinations
+    _hand_indices = list(xunique_combinations(range(7), 5))
+    """List of unique indices selecting 5 cards from a hand of 7.
+    """
+
+    def __init__(self, hand):
+        self.full_hand = hand
+        self.hand,self.residual = self.high_hand()
+
+    def __str__(self):
+        return pp_hand(self.hand + self.residual)
+
+    def high_hand(self):
+        score = hand = residual = None
+        for indices in self._hand_indices:
+            h = FiveCardHand([self.full_hand[i] for i in indices])
+            if h.score() > score:
+                residual = [self.full_hand[i]
+                            for i in range(len(self.full_hand))
+                            if i not in indices]
+                score = h.score()
+                hand = h.hand
+        return (hand, residual)