From 3f1e8193395521f84f611160b5b8ada4eda4cb6d Mon Sep 17 00:00:00 2001
From: Katherine Mantel <kmantel@princeton.edu>
Date: Thu, 16 May 2024 01:29:18 +0000
Subject: [PATCH] utilities: cached_graph_function: expose inner function using
 lru_cache

as cached_hashable_graph_function, to allow accessing or clearing cache
---
 src/graph_scheduler/utilities.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/graph_scheduler/utilities.py b/src/graph_scheduler/utilities.py
index 748b8e4..33496dc 100644
--- a/src/graph_scheduler/utilities.py
+++ b/src/graph_scheduler/utilities.py
@@ -3,7 +3,7 @@
 import inspect
 import logging
 import weakref
-from typing import Dict, Hashable, List, Set, Union
+from typing import Callable, Dict, Hashable, List, Set, Union
 
 import networkx as nx
 
@@ -262,18 +262,19 @@ def output_graph_image(
     print(f'graph_scheduler.output_graph_image: wrote {format} to {filename}')
 
 
+@functools.lru_cache()
+def cached_hashable_graph_function(func: Callable, graph: HashableDict):
+    return func(graph)
+
+
 def cached_graph_function(func):
     """
     Decorator that can be applied to cache the results of a function
     that takes a graph dependency dict
     """
-    @functools.lru_cache()
-    def cached_orig_func(graph: HashableDict):
-        return func(graph)
-
     @functools.wraps(func)
     def graph_function_wrapper(graph: typing_graph_dependency_dict):
-        return cached_orig_func(HashableDict(graph))
+        return cached_hashable_graph_function(func, HashableDict(graph))
 
     return graph_function_wrapper