ObjSim/python/objsimpy/element_tree_object.py

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as et

class MissingChildOrAttribute(Exception):
    def __init__(self, *args, **kwargs): 
        super(MissingChildOrAttribute, self).__init__(*args, **kwargs)

class EtObject:
    def __init__(self, element=None, file_path=None):
        if element is not None:
            self.element = element
        elif file_path is not None:
            self.tree = et.parse(file_path)
            self.element = self.tree.getroot()
    
    def __getattr__(self, name):
        childs = self.element.findall(name)
        objects = [EtObject(c) for c in childs]
        if len(objects) == 0:
            if not name in self.element.attrib:
                raise MissingChildOrAttribute(name)
            return self.element.attrib[name]
        else:
            return objects
        
    def __repr__(self):
        return et.tostring(self.element)

    def __str__(self):
        return self.__repr__()
        
    def save(self, file_name):
        self.tree.write(file_name)