from xml.etree import ElementTree
from xmlstruct.xml_element import XmlElement


class Optional(XmlElement):
    """
    An optional element.
    In _build, if obj is None don't build the element.
    In _parse, if element is None don't parse the element
    """

    def __init__(self, child):
        XmlElement.__init__(self, child.tag, child.attrib)
        self.child = child

    def _build(self, obj):
        if obj is None:
            return ElementTree.Element(None)
        else:
            return self.child._build(obj)

    def _parse(self, element):
        if element is None:
            return None
        else:
            return self.child._parse(element)


class Switch(XmlElement):
    """
    A conditional branch element. Switch will choose the right element according
    to a the value of keyfunc (which can be function or value).
    If default is specified, lack of match between keyfunc and cases will
    cause choosing default. If default is not specified, the result will be raising
    an Exception
    """

    def __init__(self, tag, keyfunc, cases, default=None, **kwargs):
        XmlElement.__init__(self, tag, kwargs.get("attrib"))
        self.keyfunc = keyfunc
        self.cases = cases
        self.default = default

@@ -72,16 +72,140 @@
 Element%0A
+from xmlstruct.container import ValueContainer%0Afrom xmlstruct.exceptions import SwitchNoMatchError, SwitchSeveralMatchError%0A
 %0A%0Aclass 
@@ -1182,21 +1182,16 @@
 __(self,
- tag,
  keyfunc
@@ -1258,19 +1258,20 @@
 _(self, 
-tag
+None
 , kwargs
@@ -1375,8 +1375,1810 @@
 default%0A
+%0A    def get_key(self, obj):%0A        return self.keyfunc(obj) if callable(self.keyfunc) else self.keyfunc%0A%0A    def all_cases(self):%0A        if self.default is None:%0A            return self.cases.values()%0A        else:%0A            return self.cases.values() + %5Bself.default%5D%0A%0A    def get_tag(self, parent):%0A        case_tag = %5B%5D%0A        for case in self.all_cases():%0A            if type(parent) == ElementTree.Element:%0A                # for parsing%0A                value_in_parent = parent.find(case.get_tag(parent))%0A            else:%0A                value_in_parent = parent.get(case.get_tag(parent))%0A            if value_in_parent is not None:%0A                case_tag.append(case.get_tag(parent))%0A        if len(case_tag) == 0:%0A            raise SwitchNoMatchError()%0A        elif len(case_tag) == 1:%0A            return case_tag%5B0%5D%0A        else:%0A            raise SwitchSeveralMatchError()%0A%0A    def _build(self, obj):%0A        case = self.cases.get(self.get_key(obj), self.default)%0A        if case is not None:%0A            return case._build(obj)%0A        else:%0A            raise SwitchNoMatchError()%0A%0A    def _parse(self, element):%0A        possible_objs = %5B%5D%0A        for case in self.all_cases():%0A            try:%0A                obj = case._parse(element)%0A                if type(obj) == ValueContainer:%0A                    obj = obj.value%0A                if self.cases.get(self.get_key(obj), self.default) == case:%0A                    if ElementTree.tostring(element) == case.build(obj):%0A                        possible_objs.append(obj)%0A            except Exception:%0A                continue%0A        if len(possible_objs) == 0:%0A            raise SwitchNoMatchError()%0A        elif len(possible_objs) == 1:%0A            return possible_objs%5B0%5D%0A        else:%0A            raise SwitchSeveralMatchError()%0A

Switch added