Combinatorial puzzle solver

#!/usr/local/bin/python

import re, itertools as it, sys

__all__= "LogicPuzzle",

class NotYet(Exception):
    "Raised when a rule cannot be tested because some of its items are not yet in the matrix"
    pass

class PyzzleSyntaxError(SyntaxError):
    "Raised when there is a syntax error in the description file"
    pass

class _Item(object):
    __slots__= 'name', 'column', 'category', 'row', 'value', 'rules'
    reIsNumeric= re.compile('^_\d+$')

    def __init__(self, category, name, value=None):
        self.category= category
        self.name= name
        ## self.row= self.column= None
        if value is None and self.reIsNumeric.match(name):
            self.value= int(name[1:])
            ## print name, self.value
        elif value is not None:
            self.value= value
            # deliberately we do not set the value if not provided with one

        self.rules= []

    def __eq__(self, other):
        "True when operands are in the same column"
        return self.row == other.row

    def __ne__(self, other):
        "True when operands are on different columns"
        return self.row != other.row

    def __str__(self):
        return self.name

    def _get_unknown(self):
        return not hasattr(self, 'row')
    unknown= property(_get_unknown)

class _Rule(object):
    __slots__= 'source', 'code', 'namespace',

    # here let's make the program a little version independent
    # what is reported when we don't index a list with an integer?
    _alist= [1]
    try:
        _alist['hello']
    except TypeError, exc:
        LIST_INDEX_TYPE_ERROR_ARGS= exc.args
    del _alist
    TYPE_ERROR_ARGS= LIST_INDEX_TYPE_ERROR_ARGS,

    def __init__(self, source, namespace):
        self.source= source
        self.code= compile(source, "", "eval")
        self.namespace= namespace

    def __nonzero__(self):
        try:
            return eval(self.code, self.namespace)
        except AttributeError, exc:
            if exc.args[0] in ('row','value'):
                raise NotYet
            else:
                raise AttributeError, (self.source,) + exc.args
        except NameError, exc:
            raise PyzzleSyntaxError, "rule '%s': %s" % (self.source, exc.args)

class _Category(object):
    "A category of items in the puzzle"
    __slots__= 'name', 'items', 'max_width', 'null_item'
    # name: the name of the category
    # items: list of _Item
    # max_width: used to pretty print the result
    reItem= re.compile("(?P<item>\w+)\s*(?:=\s*(?P<value>\d+))?")

    def __init__(self, name, items_as_string, custom_item_class):
        self.name= name
        self.max_width= len(name)
        self.items= []
        items= self.reItem.findall(items_as_string)
        self.null_item= custom_item_class(name, '%s' % name)
        ## print repr(items)
        for item_name, item_value in items:
            self.max_width= max(self.max_width, len(item_name))
            if item_value=='':
                item_value= None
            else:
                item_value= int(item_value)
            self.items.append(custom_item_class(self.name, item_name, item_value))

    def __iter__(self):
        for item in self.items:
            if item.unknown:
                yield item

class LogicPuzzle(object):
    "A logic puzzle as described in a list of strings (file.xreadlines comes to mind :)"
    __slots__= 'categories', 'matrix', 'items', 'category_column', 'rules', 'empty_item'

    def __init__(self, text_lines=None):
        if text_lines is not None:
            self.from_lines(text_lines)

    def _reset(self):
        "Initialise the data"
        self.categories= []
        self.matrix= []
        self.rules= []
        self.items= {'__builtins__': {}} # security precaution

    def from_lines(self, text_lines):
        "Read the puzzle description from a list of lines (or a text file object)"
        reCategory= re.compile("^(?P<category>\w+)\s*:\s*(?P<items>(?:\w+\s*(?:=\s*\d+\s*)?,?\s*)+)$")
        self._reset()
        reading_stage= "categories"
        line_count= 0
        categories= []
        for text in text_lines:
            text= text.strip()
            line_count+= 1
            if not text or text.startswith("#"): continue # skip over empty lines and comments
            if reading_stage=="categories":
                if text.startswith("--"): # we reached the end of categories
                    reading_stage= "rules"
                    item_class, self.category_column= self._create_custom_item_class(categories, self.matrix)
                    self._setup_categories(categories, item_class)
                    continue
                match= reCategory.match(text) # it MUST be a category
                if match is None: # oops!
                    raise PyzzleSyntaxError, "line %d, '%s' is not a category line" % (line_count, text)
                categories.append( (match.group("category"), match.group("items")) )
            elif reading_stage=="rules":
                new_rule= _Rule(text, self.items)
                self.rules.append(new_rule)
                # FIXME rules hung on specific items here
                for name in new_rule.code.co_names:
                    if name in self.items:
                        self.items[name].rules.append(new_rule)
            else: # normally we never come here
                raise RuntimeError, "line %d: invalid reading_stage '%s'" % (line_count, reading_stage)

    def _create_custom_item_class(self, categories, matrix):
        "internal: return a custom Item class for this puzzle and a category to column index map"
        category_columns= {}
        def _make_item_lookup_through_category(category_index, matrix):
            def _how_can_i_name_this(self):
                return matrix[self.row][category_index].value
            return _how_can_i_name_this
        for column, category_stuff in enumerate(categories):
            category_name, category_items= category_stuff
            category_columns[category_name]= column
        class CustomItem(_Item):
            pass
        for category_column, category_stuff in enumerate(categories):
            category_name, category_items= category_stuff
            setattr(CustomItem, category_name,
            property(_make_item_lookup_through_category(category_column, matrix)))
        return CustomItem, category_columns

    def _setup_categories(self, categories, item_class):
        count_of_items= None
        for category_name, category_items in categories:
            new_category= _Category(category_name, category_items, item_class)

            # assure all categories have as many items as the first
            if count_of_items is None: # this is the first category
                count_of_items= len(new_category.items)
            else:
                if len(new_category.items)!=count_of_items:
                    raise PyzzleSyntaxError, "category %s has %d items instead of %d" % \
                    (category_name, len(new_category.items), count_of_items)

            self.categories.append(new_category)
            for item in new_category.items:
                self.items[item.name]= item

    def printit(self):
        format= ' '.join(["%2d"] + ["%%-%ds" % category.max_width for category in self.categories])
        print format % tuple([0] + [category.name for category in self.categories])
        for row in xrange(len(self.categories[0].items)):
            print format % tuple([row+1] + self.matrix[row])
        print

    def solve(self):
        matrix= self.matrix
        del matrix[:] # clean the data
        row_count= len(self.categories[0].items)
        column_count= len(self.categories)
        stack= []
        for row in xrange(row_count):
            self.matrix.append([])
            for column, category in enumerate(self.categories):
                if column==0:
                    item= category.items[row]
                    self.matrix[row].append(item) # first column does not change
                    item.row= row
                    item.column= 0
                else:
                    self.matrix[row].append(category.null_item)
        def recurse(row, column):
        ## try:
        ## if self.rules[0]:
        ## self.printit()
        ## print "here we are!"
        ## print '\n'.join([rule.source for rule in self.rules])
        ## raw_input()
        ## except NotYet:
        ## pass
            if column>=column_count:
                recurse(row+1, 1)
                return
            if row>=row_count:
            ## _= self.items['LeadPiping']
            ## print dir(_)
            ## print "name", _.name
            ## print "value", _.value
            ## print "row", _.row
            ## print "column", _.column
            ## print _.DiscoveryDate
            ## self.printit()
            ## print self.items['LeadPiping'].DiscoveryDate
            ## print self.items['RobertTrench'].DiscoveryDate
            ## sys.exit(0)
            # check the rules
                all_rules_ok= True
                for rule in self.rules:
                    try:
                        if not rule:
                        ## print "failed:", rule.source, " "
                        ## raw_input()
                            all_rules_ok= False
                            break
                    except NotYet: # here ???
                        print rule.source, self.items.keys()
                        for name in rule.code.co_names:
                            if name in self.items:
                                print name, "=", self.items[name].row, self.items[name].column
                        sys.exit(1)
                if all_rules_ok:
                    self.printit()
                    print "******** solution! ************"
                    print '\n'.join(['%s: %s' % (rule.source, bool(rule)) for rule in self.rules])
                return
            for item in self.categories[column]:
                matrix[row][column]= item
                item.row, item.column= row, column
                for rule in item.rules: # item.rules
                    try:
                        if not rule:
                        ## self.printit()
                        ## print rule.source, " "
                        ## raw_input()
                            break
                    except NotYet:
                        pass
                else:
                    recurse(row, column+1)
                del item.row, item.column
                matrix[row][column]= self.categories[column].null_item
        recurse(0, 1)

if __name__=='__main__':
    p = LogicPuzzle(file(sys.argv[1]))
    p.solve()