Notify subclasses of Overlay which <source> is selected

1 files changed, 10 insertions, 6 deletions

diff --git a/layman/overlays/overlay.py b/layman/overlays/overlay.py
index aa20f9d..e613ec2 100644
--- a/layman/overlays/overlay.py
+++ b/layman/overlays/overlay.py
@@ -83,12 +83,13 @@ class Overlay(object):
             raise Exception('Overlay is missing a "name" entry!')
 
         _source = xml.find('source')
-        if _source != None:
-            self.src = ensure_unicode(_source.text.strip())
-        elif 'src' in xml.attrib:
-            self.src = ensure_unicode(xml.attrib['src'])
-        else:
-            raise Exception('Overlay "' + self.name + '" is missing a "source" entry!')
+        if _source == None:
+            if 'src' in xml.attrib:
+                _source = ET.Element('source')
+                _source.text = xml.attrib['src']
+            else:
+                raise Exception('Overlay "' + self.name + '" is missing a "source" entry!')
+        self._set_source(_source)
 
         _owner = xml.find('owner')
         if _owner == None:
@@ -157,6 +158,9 @@ class Overlay(object):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def _set_source(self, source_elem):
+        self.src = ensure_unicode(source_elem.text.strip())
+
     def set_priority(self, priority):
         '''Set the priority of this overlay.'''