[common] Improve handling of version/fork setup

beat/cmdline/common.py

@@ -151,35 +151,32 @@ class Selector(object):
         self.prefix = prefix  # the root of the directory
         self.path = os.path.join(self.prefix, ".beat", "selected.json")
 
-        status = False
+        self.__version = {}
+        self.__fork = {}
+        self.__versionables = [
+            "dataformat",
+            "database",
+            "algorithm",
+            "library",
+            "toolchain",
+            "experiment",
+            "protocoltemplate",
+        ]
+
+        self.__forkables = [
+            "dataformat",
+            "database",
+            "algorithm",
+            "library",
+            "toolchain",
+            "experiment",
+            "protocoltemplate",
+        ]
+
         if os.path.exists(self.path):
-            status = self.load()
-
-        if not status:
-            self.__version = {
-                "dataformat": dict(),
-                "database": dict(),
-                "algorithm": dict(),
-                "library": dict(),
-                "toolchain": dict(),
-                "experiment": dict(),
-                "protocoltemplate": dict(),
-            }
-            self.__fork = {
-                "dataformat": dict(),
-                "database": dict(),
-                "algorithm": dict(),
-                "library": dict(),
-                "toolchain": dict(),
-                "experiment": dict(),
-                "protocoltemplate": dict(),
-            }
+            self.load()
         else:
-            for type_ in TYPE_GLOB.keys():
-                if type_ not in self.__version:
-                    self.__version[type_] = dict()
-                if type_ not in self.__fork:
-                    self.__fork[type_] = dict()
+            self.__ensure_entries()
 
     def __enter__(self):
         """Implements our context manager"""
@@ -189,6 +186,17 @@ class Selector(object):
         """Implements our context manager"""
         self.save()
 
+    def __ensure_entries(self):
+        """Ensure all types have an entry"""
+
+        for type_ in self.__versionables:
+            if type_ not in self.__version:
+                self.__version[type_] = dict()
+
+        for type_ in self.__forkables:
+            if type_ not in self.__fork:
+                self.__fork[type_] = dict()
+
     def fork(self, type, src, dst):
         """Registers that object ``dst`` is a fork of object ``src``"""
 
@@ -248,6 +256,8 @@ class Selector(object):
 
         self.__fork = data["fork"]
         self.__version = data["version"]
+        self.__ensure_entries()
+
         return True
 
     def save(self):