Add a helper method to Image to parse ancestor string.

util/imagetree.py

@@ -1,3 +1,6 @@
+from collections import defaultdict
+
+
 class ImageTreeNode(object):
   """ A node in the image tree. """
   def __init__(self, image, child_map):
@@ -9,7 +12,7 @@ class ImageTreeNode(object):
 
   @property
   def children(self):
-    return self._child_map.get(str(self.image.id), [])
+    return self._child_map[self.image.id]
 
   def add_tag(self, tag):
     self.tags.append(tag)
@@ -20,18 +23,18 @@ class ImageTree(object):
 
   def __init__(self, all_images, all_tags, base_filter=None):
     self._image_map = {}
-    self._child_map = {}
+    self._child_map = defaultdict(list)
 
     self._build(all_images, all_tags, base_filter)
 
-  def _build(self,  all_images, all_tags, base_filter=None):
+  def _build(self, all_images, all_tags, base_filter=None):
     # Build nodes for each of the images.
     for image in all_images:
-      ancestors = image.ancestors.split('/')[1:-1]
+      ancestors = image.ancestor_id_list()
 
       # Filter any unneeded images.
       if base_filter is not None:
-        if image.id != base_filter and not str(base_filter) in ancestors:
+        if image.id != base_filter and not base_filter in ancestors:
           continue
 
       # Create the node for the image.
@@ -39,11 +42,8 @@ class ImageTree(object):
       self._image_map[image.id] = image_node
 
       # Add the node to the child map for its parent image (if any).
-      parent_image_id = image.ancestors.split('/')[-2] if image.ancestors else None
-      if parent_image_id:
-        if not parent_image_id in self._child_map:
-          self._child_map[parent_image_id] = []
-
+      parent_image_id = image.parent_id
+      if parent_image_id is not None:
         self._child_map[parent_image_id].append(image_node)
 
     # Build the tag map.
@@ -54,7 +54,6 @@ class ImageTree(object):
 
       image_node.add_tag(tag.name)
 
-
   def find_longest_path(self, image_id, checker):
     """ Returns a list of images representing the longest path that matches the given
         checker function, starting from the given image_id *exclusive*.
@@ -65,7 +64,6 @@ class ImageTree(object):
 
     return self._find_longest_path(start_node, checker, -1)[1:]
 
-
   def _find_longest_path(self, image_node, checker, index):
     found_path = []
 
@@ -79,7 +77,6 @@ class ImageTree(object):
 
     return [image_node.image] + found_path
 
-
   def tag_containing_image(self, image):
     """ Returns the name of the closest tag containing the given image. """
     if not image:
@@ -99,4 +96,4 @@ class ImageTree(object):
       if found is not None:
         return found
 
-    return None
+    return None