X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;ds=sidebyside;f=collect%2Fshared_dict.py;h=e0a42f2c55c2fc865b0c89642d458ae26009c224;hb=d2357ff35e7752ae3eb6caa2813c35c17fea778b;hp=7c84c14c073743ea1e452d58393f8f29d280ed23;hpb=6ba90a1f30f1c0cf4df12fcd0c62181f29bc3668;p=python_utils.git

diff --git a/collect/shared_dict.py b/collect/shared_dict.py
index 7c84c14..e0a42f2 100644
--- a/collect/shared_dict.py
+++ b/collect/shared_dict.py
@@ -31,16 +31,7 @@ import pickle
 from contextlib import contextmanager
 from functools import wraps
 from multiprocessing import RLock, shared_memory
-from typing import (
-    Any,
-    Dict,
-    Generator,
-    ItemsView,
-    Iterator,
-    KeysView,
-    Optional,
-    ValuesView,
-)
+from typing import Any, Dict, Generator, ItemsView, Iterator, KeysView, Optional, ValuesView
 
 from decorator_utils import synchronized
 
@@ -74,6 +65,7 @@ class SharedDict(object):
         super().__init__()
         self.name = name
         self._serializer = PickleSerializer()
+        assert size_bytes is None or size_bytes > 0
         self.shared_memory = self._get_or_create_memory_block(name, size_bytes)
         self._ensure_memory_initialization()
         self.lock = RLock()
@@ -89,13 +81,11 @@ class SharedDict(object):
         try:
             return shared_memory.SharedMemory(name=name)
         except FileNotFoundError:
-            assert size_bytes
+            assert size_bytes is not None
             return shared_memory.SharedMemory(name=name, create=True, size=size_bytes)
 
     def _ensure_memory_initialization(self):
-        memory_is_empty = (
-            bytes(self.shared_memory.buf).split(SharedDict.NULL_BYTE, 1)[0] == b''
-        )
+        memory_is_empty = bytes(self.shared_memory.buf).split(SharedDict.NULL_BYTE, 1)[0] == b''
         if memory_is_empty:
             self.clear()
 
@@ -107,7 +97,8 @@ class SharedDict(object):
     def cleanup(self) -> None:
         if not hasattr(self, 'shared_memory'):
             return
-        self.shared_memory.unlink()
+        with SharedDict.MPLOCK:
+            self.shared_memory.unlink()
 
     def clear(self) -> None:
         self._save_memory({})