Fix cache_restoring context

python/dune/codegen/generation/context.py

@@ -37,16 +37,17 @@ class _CacheRestoringContext(object):
     def __enter__(self):
         from dune.codegen.generation.cache import _generators as g
         self.cache = {}
-        for i in g:
-            self.cache[i] = {}
-            for k, v in i._memoize_cache.items():
-                self.cache[i][k] = v
+        for original_func, cache_func in g.items():
+            self.cache[original_func] = {}
+            for k, v in cache_func._memoize_cache.items():
+                self.cache[original_func][k] = v
 
     def __exit__(self, exc_type, exc_value, traceback):
+        from dune.codegen.generation.cache import _generators as g
         for i, c in self.cache.items():
-            i._memoize_cache = {}
+            g[i]._memoize_cache = {}
             for k, v in c.items():
-                i._memoize_cache[k] = v
+                g[i]._memoize_cache[k] = v
 
 
 def cache_restoring():