Sudoku solver

sudoku.py

def valid(puzzle):
    m = [list(row) for row in puzzle.split('\n')]
    for row in m:
        counts = set()
        for ch in row:
            if ch in counts and '_' != ch:
                return False
            counts.add(ch)
    for column in zip(*m):
        counts = set()
        for ch in column:
            if ch in counts and '_' != ch:
                return False
            counts.add(ch)
    for i in range(3):
        for j in range(3):
            counts = set()
            for ch in (
                x
                for row in m[i * 3:(i + 1) * 3]
                for x in row[j * 3:(j + 1) * 3]
            ):
                if ch in counts and '_' != ch:
                    return False
                counts.add(ch)
    return True


def solve(puzzle):
    s = [puzzle]
    while True:
        p = s.pop()
        if valid(p):
            if '_' in p:
                for i in range(1, 10):
                    s.append(p.replace('_', str(i), 1))
            else:
                return p
                
                
def format(puzzle):
  def format_row(row):
    return ' | '.join((row[:3], row[3:6], row[6:]))
  fmt = []
  for i, row in enumerate(puzzle.split('\n')):
    if i > 0 and i % 3 == 0:
      fmt.append('---------------')
    fmt.append(format_row(row))
  return '\n'.join(fmt)


# Example
if __name__ == '__main__':
    puzzle = """__729___5
    ____84_26
    _1_______
    1____8_37
    __6_7_8__
    73_4____2
    _______9_
    38_96____
    2___576__"""
    print(format(solve(puzzle)))

sudoku_test.py

import unittest

from sudoku import solve


class SudokuTests(unittest.TestCase):
    def test_one_blank(self):
        puzzle = '\n'.join([
            '_68579231',
            '973812465',
            '215634978',
            '897423516',
            '546981723',
            '132756849',
            '681297354',
            '754368192',
            '329145687',
        ])
        expected = '\n'.join([
            '468579231',
            '973812465',
            '215634978',
            '897423516',
            '546981723',
            '132756849',
            '681297354',
            '754368192',
            '329145687',
        ])
        self.assertEqual(solve(puzzle), expected)

    def test_medium(self):
        puzzle = '\n'.join([
            '____79___',
            '9__8_2_6_',
            '_15__4_7_',
            '__74___1_',
            '__69817__',
            '_3___68__',
            '_8_2__35_',
            '_5_3_8__2',
            '___14____',
        ])
        expected = '\n'.join([
            '468579231',
            '973812465',
            '215634978',
            '897423516',
            '546981723',
            '132756849',
            '681297354',
            '754368192',
            '329145687',
        ])
        self.assertEqual(solve(puzzle), expected)

    def test_hard(self):
        puzzle = '\n'.join([
            '__729___5',
            '____84_26',
            '_1_______',
            '1____8_37',
            '__6_7_8__',
            '73_4____2',
            '_______9_',
            '38_96____',
            '2___576__',
        ])
        expected = '\n'.join([
            '467291385',
            '953784126',
            '812536479',
            '194628537',
            '526379841',
            '738415962',
            '675143298',
            '381962754',
            '249857613',
        ])
        self.assertEqual(solve(puzzle), expected)


if __name__ == '__main__':
    unittest.main()