images.py

from collections import OrderedDict, namedtuple
from copy import deepcopy
from typing import Callable, List, Optional

Node = namedtuple('Node', 'name, branches')

Image = namedtuple('Image', 'name, build_from')

images = [
    Image('image1', 'centos:7'),  # root
    Image('image7', 'image1'),
    Image('image2', 'image1'),
    Image('image3', 'image2'),
    Image('image4', 'image2'),
    Image('image6', 'image4'),
    Image('image8', 'image7'),
    Image('image5', 'ubuntu:18.04'),  # root
    Image('image10', 'image5'),
    Image('image11', 'image5'),
]


def get_images_to_rebuild(all_images: List[Image],
                          changed_images: List[str]) -> List[str]:
    trees = make_trees(all_images)
    changed: List[str] = []
    for tree in trees:
        for image in changed_images:
            cut = cut_tree(tree, image)
            changed += flatten_tree(cut)
    return uniq(compact(changed) + changed_images)


def uniq(lst: List) -> List:
    """Uniquify list `lst` preserving original order"""
    return list(OrderedDict.fromkeys(lst))


def make_trees(images: List[Image]) -> List[Node]:
    """Build trees out of images.

    Each tree is an image dependency graph.
        Note: circular dependencies are not supported

    In the root of a tree is an image which `.build_from` is not in the image
    names list.
    """

    image_names = [img.name for img in images]
    build_from = [img.build_from for img in images]

    roots = [
        build_from_name for build_from_name in build_from
        if build_from_name not in image_names
    ]

    trees = []
    for root_node in roots:
        tree = Node(root_node, [])
        trees.append(tree)

        def _build_tree(tree, images: List[Image]):
            for image in images:
                if image.build_from == tree.name:
                    node = Node(image.name, [])
                    tree.branches.append(node)
                    _build_tree(node, images)

        _build_tree(tree, images)
    return trees


def compact(lst: List) -> List:
    return [value for value in lst if value]


def traverse(tree: Node, cb: Callable[[Node], None]) -> None:
    """Walk (depth first) a tree and apply `cb` function to each node."""

    tree = deepcopy(tree)

    def _traverse(tree):
        for node in tree.branches:
            cb(node)
            _traverse(node)

    _traverse(tree)


def flatten_tree(tree: Optional[Node]) -> List[str]:
    """Traverse a tree and return all node names.

    Given a tree:

             (n0)
            /  |  \
          n1   n2  n3
              / \
             n5 n6

    return [n0, n1, n2, n5, n6, n3]
    """

    if not tree:
        return []

    names = [tree.name]

    def collect_name(node):
        names.append(node.name)

    traverse(tree, collect_name)
    return names


def cut_tree(tree: Node, cut_to: str) -> Optional[Node]:
    """Cut tree returning a new tree starting from `cut_to`.

    Given a tree:

             (n0)
            /  |  \
          n1   n2  n3
          /   / \   \
        n4   n5 n6   n7

    And a node name (cut_to) n2, return:
            n2
           / \
          n5 n6
    """

    tree = deepcopy(tree)
    if tree.name == cut_to:
        return tree
    found = None
    for node in tree.branches:
        found = cut_tree(node, cut_to)
        if found:
            return found
    return found


def test_get_images_to_rebuild():
    test_cases = [
        [['image3', 'image2'], ['image2', 'image3', 'image4', 'image6']],
        [['image1', 'image4'], [
            'image1', 'image7', 'image8', 'image2', 'image3', 'image4',
            'image6'
        ]],
        [['image4'], ['image4', 'image6']],
        [['image6'], ['image6']],
        [[], []],
    ]
    for changed, answer in test_cases:
        result = get_images_to_rebuild(images, changed)
        print(f"  changed: {changed}\n\tanswer: {answer}\n\tresult: {result}")
        assert result == answer
        print()
    print('All tests pass: ' + len(test_cases) * '🍏')