From 856698ee9c745dad6181dc8a78d3d2c0a34abd50 Mon Sep 17 00:00:00 2001 From: Christian Kellner Date: Thu, 5 Mar 2020 18:49:55 +0100 Subject: [PATCH] objectstore: keep track of created objects Keep track of all created objects via weak references. Add support to use ObjectStore as context manager and ensure that all objects are cleaned up when the context is exited. --- osbuild/objectstore.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/osbuild/objectstore.py b/osbuild/objectstore.py index d4312285..5f71da58 100644 --- a/osbuild/objectstore.py +++ b/osbuild/objectstore.py @@ -5,6 +5,7 @@ import hashlib import os import subprocess import tempfile +import weakref from typing import Optional from . import treesum @@ -212,6 +213,7 @@ class ObjectStore: os.makedirs(self.objects, exist_ok=True) os.makedirs(self.refs, exist_ok=True) os.makedirs(self.tmp, exist_ok=True) + self._objs = weakref.WeakSet() def contains(self, object_id): if not object_id: @@ -257,6 +259,8 @@ class ObjectStore: # copying of the data takes places at this point obj.base = base_id + self._objs.add(obj) + return obj def commit(self, obj: Object, object_id: str) -> str: @@ -298,3 +302,15 @@ class ObjectStore: obj.base = object_id return treesum_hash + + def cleanup(self): + """Cleanup all created Objects that are still alive""" + for obj in self._objs: + obj.cleanup() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.cleanup() + return exc_type is None