From 5a20ccf2f727b361a62c81e7d3c6be4a6bdf7bae Mon Sep 17 00:00:00 2001
From: jpic <jpic@yourlabs.org>
Date: Sat, 8 Aug 2020 16:08:01 +0200
Subject: [PATCH] Layers refactor

---
 shlax/image.py           | 40 +++++++++++++++++++++++++++++++++-------
 shlax/targets/buildah.py | 32 ++++++++------------------------
 2 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/shlax/image.py b/shlax/image.py
index f21737b..33c80e3 100644
--- a/shlax/image.py
+++ b/shlax/image.py
@@ -1,21 +1,52 @@
-import copy
+import json
 import os
 import re
 
 
+class Layers(set):
+    def __init__(self, image):
+        self.image = image
+
+    async def ls(self, target):
+        """Fetch layers from localhost"""
+        ret = set()
+        results = await target.parent.exec(
+            'buildah images --json',
+            quiet=True,
+        )
+        results = json.loads(results.out)
+
+        prefix = 'localhost/' + self.image.repository + ':layer-'
+        for result in results:
+            if not result.get('names', None):
+                continue
+            for name in result['names']:
+                if name.startswith(prefix):
+                    self.add(name)
+        return self
+
+    async def rm(self, target, tags=None):
+        """Drop layers for this image"""
+        if tags is None:
+            tags = [layer for layer in await self.ls(target)]
+        await target.exec('podman', 'rmi', *tags)
+
+
 class Image:
     PATTERN = re.compile(
         '^((?P<backend>[a-z]*)://)?((?P<registry>[^/]*[.][^/]*)/)?((?P<repository>[^:]+))?(:(?P<tags>.*))?$'  # noqa
         , re.I
     )
 
-    def __init__(self, arg=None, format=None, backend=None, registry=None, repository=None, tags=None):
+    def __init__(self, arg=None, format=None, backend=None, registry=None,
+                 repository=None, tags=None):
         self.arg = arg
         self.format = format
         self.backend = backend
         self.registry = registry
         self.repository = repository
         self.tags = tags or []
+        self.layers = Layers(self)
 
         match = re.match(self.PATTERN, arg)
         if match:
@@ -53,8 +84,3 @@ class Image:
 
         for tag in self.tags:
             await action.exec('buildah', 'push', f'{self.repository}:{tag}')
-
-    def layer(self, key):
-        layer = copy.deepcopy(self)
-        layer.tags = ['layer-' + key]
-        return layer
diff --git a/shlax/targets/buildah.py b/shlax/targets/buildah.py
index 2652386..5314a61 100644
--- a/shlax/targets/buildah.py
+++ b/shlax/targets/buildah.py
@@ -43,16 +43,15 @@ class Buildah(Target):
 
         if not self.is_runnable():
             os.execvp('buildah', ['buildah', 'unshare'] + sys.argv)
-            # program has been replaced
+            return  # process has been replaced
 
-        layers = await self.layers()
-        keep = await self.cache_setup(layers, *actions)
+        layers = await self.image.layers.ls(self)
+        keep = await self.cache_setup(self.image.layers, *actions)
         keepnames = [*map(lambda x: 'localhost/' + str(x), keep)]
-        self.invalidate = [name for name in layers if name not in keepnames]
+        self.invalidate = [name for name in self.image.layers if name not in keepnames]
         if self.invalidate:
             self.output.info('Invalidating old layers')
-            await self.parent.exec(
-                'buildah', 'rmi', *self.invalidate, raises=False)
+            await self.image.layers.rm(self.parent, self.invalidate)
 
         if actions:
             actions = actions[len(keep):]
@@ -68,23 +67,6 @@ class Buildah(Target):
 
         return await super().__call__(*actions)
 
-    async def layers(self):
-        ret = set()
-        results = await self.parent.exec(
-            'buildah images --json',
-            quiet=True,
-        )
-        results = json.loads(results.out)
-
-        prefix = 'localhost/' + self.image.repository + ':layer-'
-        for result in results:
-            if not result.get('names', None):
-                continue
-            for name in result['names']:
-                if name.startswith(prefix):
-                    ret.add(name)
-        return ret
-
     async def cache_setup(self, layers, *actions):
         keep = []
         self.image_previous = Image(self.base)
@@ -115,7 +97,9 @@ class Buildah(Target):
             action_key = str(action)
         key = prefix + action_key
         sha1 = hashlib.sha1(key.encode('ascii'))
-        return self.image.layer(sha1.hexdigest())
+        action_image = copy.deepcopy(self.image)
+        action_image.tags = ['layer-' + sha1.hexdigest()]
+        return action_image
 
     async def action(self, action, reraise=False):
         stop = await super().action(action, reraise)