refactor: robust dictionary access

ing-bank · Jun 14, 2022 · f438394 · f438394
1 parent 17ac6d8
commit f438394
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/popmon/analysis/profiling/profiles.py b/popmon/analysis/profiling/profiles.py
@@ -22,7 +22,7 @@
 
 class Profiles:
     _profile_descriptions = {}
-    _profile_funcs = {-1: {}, 1: {}, 2: {}}
+    _profile_funcs = {-1: {}}
 
     @classmethod
     def register(
@@ -46,6 +46,9 @@ def f(func: Callable):
                     cls._profile_descriptions[k] = d
             else:
                 cls._profile_descriptions[key] = description
+
+            if dim not in cls._profile_funcs:
+                cls._profile_funcs[dim] = {}
             cls._profile_funcs[dim][(key, htype)] = func
             return func
 
@@ -65,7 +68,7 @@ def merge(d1, d2):
         if dim is None:
             v = cls._profile_funcs[-1]
         else:
-            v = merge(cls._profile_funcs[dim], cls._profile_funcs[-1])
+            v = merge(cls._profile_funcs.get(dim, {}), cls._profile_funcs[-1])
 
         return v