Composite - GoF

composite.py

class Component:
  def draw(self):
    raise NotImplementedError()

  def add(self, component):
    raise ComponentError()
  
  def remove(self, component):
    raise ComponentError()

  def getChild(self, i):
    raise ComponentError() 

class Leaf(Component):
  def draw(self):
    print 'Leaf'

class CompositeComponent(Component):
  
  def __init__(self):
      self.children = []
  
  def draw(self):
    print 'Composite element. Contains of: '
    for child in self.children:
      child.draw()
   
  def add(self, element):
    self.children.append(element)

  def remove(self):
    if len(self.children) != 0:
      self.children.remove(len(self.children) - 1)

  def get(self, i):
    if i < len(self.children):
      return self.children(i)

class ComponentError(Exception):
  pass

leaf1 = Leaf()
root = CompositeComponent()
leaf2 = Leaf()
root.add(leaf1)

leaf2 = Leaf()
leaf3 = Leaf()
cc2 = CompositeComponent()
cc2.add(leaf2)
cc2.add(leaf3)

root.add(cc2)

root.draw()