From d66e844c1af20c619311dee124232f859fcded99 Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin" <git_commits@alexgoodwin.dev>
Date: Thu, 19 Sep 2024 17:04:54 +0900
Subject: [PATCH] use a 'with' context

---
 folder_paths.py |  8 ++++++++
 server.py       | 20 +++++++++-----------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/folder_paths.py b/folder_paths.py
index 40d3f1875..0ad8bc184 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -66,6 +66,14 @@ class CacheHelper:
     def clear(self):
         self.cache.clear()
 
+    def __enter__(self):
+        self.active = True
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.active = False
+        self.clear()
+
 cache_helper = CacheHelper()
 
 extension_mimetypes_cache = {
diff --git a/server.py b/server.py
index 86919b846..9321e4d08 100644
--- a/server.py
+++ b/server.py
@@ -552,17 +552,15 @@ class PromptServer():
 
         @routes.get("/object_info")
         async def get_object_info(request):
-            folder_paths.cache_helper.clear()
-            folder_paths.cache_helper.active = True
-            out = {}
-            for x in nodes.NODE_CLASS_MAPPINGS:
-                try:
-                    out[x] = node_info(x)
-                except Exception as e:
-                    logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
-                    logging.error(traceback.format_exc())
-            folder_paths.cache_helper.active = False
-            return web.json_response(out)
+            with folder_paths.cache_helper:
+                out = {}
+                for x in nodes.NODE_CLASS_MAPPINGS:
+                    try:
+                        out[x] = node_info(x)
+                    except Exception as e:
+                        logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
+                        logging.error(traceback.format_exc())
+                return web.json_response(out)
 
         @routes.get("/object_info/{node_class}")
         async def get_object_info_node(request):