fix: support **kwargs in taskgraph.util.memoize

Previously this decorator only accepted `*args`, so any memoized
function that used `**kwargs` could no longer be called via keyword
arguments.

This is handled by replacing our custom memoize implementation with the
built-in `functools.lru_cache(maxsize=None)` which is equivalent.

Author: ahal

src/taskgraph/transforms/run/__init__.py

@@ -167,7 +167,8 @@ def add_resource_monitor(config, tasks):
                 config.graph_config, task["worker-type"]
             )
             # Normalise worker os so that linux-bitbar and similar use linux tools.
-            worker_os = worker_os.split("-")[0]
+            if worker_os:
+                worker_os = worker_os.split("-")[0]
             if "win7" in task["worker-type"]:
                 arch = "32"
             else:

src/taskgraph/util/memoize.py

@@ -2,39 +2,6 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
-# Imported from
-# https://searchfox.org/mozilla-central/rev/c3ebaf6de2d481c262c04bb9657eaf76bf47e2ac/python/mozbuild/mozbuild/util.py#923-949
-
-
 import functools
 
-
-class memoize(dict):
-    """A decorator to memoize the results of function calls depending
-    on its arguments.
-    Both functions and instance methods are handled, although in the
-    instance method case, the results are cache in the instance itself.
-    """
-
-    def __init__(self, func):
-        self.func = func
-        functools.update_wrapper(self, func)
-
-    def __call__(self, *args):
-        if args not in self:
-            self[args] = self.func(*args)
-        return self[args]
-
-    def method_call(self, instance, *args):
-        name = "_%s" % self.func.__name__
-        if not hasattr(instance, name):
-            setattr(instance, name, {})
-        cache = getattr(instance, name)
-        if args not in cache:
-            cache[args] = self.func(instance, *args)
-        return cache[args]
-
-    def __get__(self, instance, cls):
-        return functools.update_wrapper(
-            functools.partial(self.method_call, instance), self.func
-        )
+memoize = functools.lru_cache(maxsize=None)  # backwards compatibility shim

test/test_util_memoize.py

@@ -150,7 +150,7 @@ def assert_artifact_refs(monkeypatch):
 
     def inner(input, output):
         # Clear memoized function
-        get_root_url.clear()
+        get_root_url.cache_clear()
         taskid_for_edge_name = {"edge%d" % n: "tid%d" % n for n in range(1, 4)}
         assert (
             resolve_task_references(

test/test_util_parameterization.py

@@ -28,30 +28,30 @@ def mock_production_taskcluster_root_url(monkeypatch):
 
 @pytest.fixture
 def root_url(mock_environ):
-    tc.get_root_url.clear()
+    tc.get_root_url.cache_clear()
     return tc.get_root_url(False)
 
 
 @pytest.fixture
 def proxy_root_url(monkeypatch, mock_environ):
     monkeypatch.setenv("TASK_ID", "123")
     monkeypatch.setenv("TASKCLUSTER_PROXY_URL", "https://taskcluster-proxy.net")
-    tc.get_root_url.clear()
+    tc.get_root_url.cache_clear()
     return tc.get_root_url(True)
 
 
 def test_get_root_url(monkeypatch):
-    tc.get_root_url.clear()
+    tc.get_root_url.cache_clear()
     assert tc.get_root_url(False) == tc.PRODUCTION_TASKCLUSTER_ROOT_URL
 
     custom_url = "https://taskcluster-root.net"
     monkeypatch.setenv("TASKCLUSTER_ROOT_URL", custom_url)
-    tc.get_root_url.clear()
+    tc.get_root_url.cache_clear()
     assert tc.get_root_url(False) == custom_url
 
     # trailing slash is normalized
     monkeypatch.setenv("TASKCLUSTER_ROOT_URL", custom_url + "/")
-    tc.get_root_url.clear()
+    tc.get_root_url.cache_clear()
     assert tc.get_root_url(False) == custom_url
 
     with pytest.raises(RuntimeError):
@@ -66,7 +66,7 @@ def test_get_root_url(monkeypatch):
     proxy_url = "https://taskcluster-proxy.net"
     monkeypatch.setenv("TASKCLUSTER_PROXY_URL", proxy_url)
     assert tc.get_root_url(True) == proxy_url
-    tc.get_root_url.clear()
+    tc.get_root_url.cache_clear()
 
     # no default set
     monkeypatch.setattr(tc, "PRODUCTION_TASKCLUSTER_ROOT_URL", None)
@@ -382,8 +382,8 @@ def test_list_task_group_incomplete_tasks(responses, root_url):
 
 
 def test_get_ancestors(responses, root_url):
-    tc.get_task_definition.clear()
-    tc._get_deps.clear()
+    tc.get_task_definition.cache_clear()
+    tc._get_deps.cache_clear()
     base_url = f"{root_url}/api/queue/v1/task"
     responses.add(
         responses.GET,
@@ -457,8 +457,8 @@ def test_get_ancestors(responses, root_url):
 
 
 def test_get_ancestors_string(responses, root_url):
-    tc.get_task_definition.clear()
-    tc._get_deps.clear()
+    tc.get_task_definition.cache_clear()
+    tc._get_deps.cache_clear()
     base_url = f"{root_url}/api/queue/v1/task"
     responses.add(
         responses.GET,