decision tree sketch

decisiontree.py

class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self


class Choices(dict):
    pass


class CannotDecide(ValueError):
    """
    Given the passed objects, the tree does not
    produce a conclusive result
    """


def _descend_object(obj, path_parts):
    if len(path_parts) == 0:
        return obj
    return _descend_object(getattr(obj, path_parts[0]), path_parts[1:])


class DecisionTree(object):
    def __init__(self, attribute, choices):
        self._attr = attribute
        self._choices = choices

    @property
    def _attr_path(self):
        return self._attr.split(".")

    def decide(self, **kwargs):
        """
        :raises: AttributeError, CannotDecide
        """
        toplevel = AttrDict(**kwargs)
        # raises: AttributeError if missing attr
        key = _descend_object(toplevel, self._attr_path)
        try:
            val = self._choices[key]
        except KeyError:
            raise CannotDecide
        if isinstance(val, DecisionTree):
            return val.decide(**kwargs)
        return val

test_decisiontree.py

from unittest import TestCase
from mock import Mock

from decisiontree import DecisionTree, Choices, CannotDecide


class DecisionTreeTestCase(TestCase):
    def test_can_choose_shallow(self):
        tree = DecisionTree(
            "a",
            Choices(**{
                1: "one",
                2: "two",
                3: "three"
            })
        )
        a = 1
        assert tree.decide(a=a) == "one"

    def test_can_choose_deep_attr(self):
        tree = DecisionTree(
            "a.b.c",
            Choices(**{
                1: "one",
                2: "two",
                3: "three"
            })
        )
        a = Mock(b=Mock(c=2))
        assert tree.decide(a=a) == "two"

    def test_can_choose_deep_tree(self):
        tree = DecisionTree(
            "a.b",
            Choices(**{
                1: DecisionTree(
                    "p.q",
                    Choices(
                        hut="wow",
                        sut="pow",
                        rut="cow",
                    )
                ),
                2: DecisionTree(
                    "p.q",
                    Choices(
                        hut="bob",
                        sut="sob",
                        rut="cob",
                    )
                )
            })
        )
        a = Mock(b=2)
        p = Mock(q="sut")
        assert tree.decide(a=a, p=p) == "sob"

    def test_missing_attr_raises(self):
        tree = DecisionTree(
            "a",
            Choices(**{
                1: "one",
                2: "two",
                3: "three"
            })
        )
        b = 1
        with self.assertRaises(AttributeError):
            # b is not a
            assert tree.decide(b=b)

    def test_no_choice_raises(self):
        tree = DecisionTree(
            "a",
            Choices(**{
                1: "one",
                2: "two",
                3: "three"
            })
        )
        a = 4
        with self.assertRaises(CannotDecide):
            assert tree.decide(a=a)