Merge pull request #1322 from linsword13/ws-cache

Add a decorator for Ramble object to cache command results

Author: douglasjacobsen

lib/ramble/ramble/workspace/workspace.py

@@ -478,6 +478,9 @@ def __init__(self, root, dry_run=False, read_default_template=True):
         # This is currently used as a cache for reading per-object template contents.
         self._inmem_file_cache = {}
 
+        # A workspace-level cache that's used by objects to cache command returns.
+        self.object_command_cache = {}
+
         self.results = self.default_results()
 
         # A cache structured as {pkg_man: {env_name: pkg_list}}.

var/ramble/repos/builtin/base_classes/object-mixin/base_class.py

@@ -5,6 +5,7 @@
 # <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
 # option. This file may not be copied, modified, or distributed
 # except according to those terms.
+import functools
 import os
 from html import escape
 from typing import List
@@ -210,3 +211,26 @@ def to_html_docs(self, out, obj_def):
 
     def _format_docs_details(self, _out):
         """Hook for objects to add extra documentation."""
+
+    @staticmethod
+    def workspace_cache(method):
+        """
+        A decorator that caches the result of a ramble object's instance method into the workspace.
+        """
+
+        @functools.wraps(method)
+        def wrapper(obj, *args, **kwargs):
+            # Precondition: the method should be called with a `workspace=` argument.
+            ws = kwargs["workspace"]
+            key = (
+                obj.origin_type,
+                obj.name,
+                method.__name__,
+                args,
+                frozenset(kwargs.items()),
+            )
+            if key not in ws.object_command_cache:
+                ws.object_command_cache[key] = method(obj, *args, **kwargs)
+            return ws.object_command_cache[key]
+
+        return wrapper

var/ramble/repos/builtin/base_classes/object-mixin/test/object-mixin.py

@@ -0,0 +1,62 @@
+# Copyright 2022-2025 The Ramble Authors
+#
+# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+# option. This file may not be copied, modified, or distributed
+# except according to those terms.
+import unittest.mock as mock
+
+import pytest
+
+from ramble.repository import get_base_class
+
+ObjectMixin = get_base_class("object-mixin")
+
+
+class TestObject(ObjectMixin):
+    """A test class for the ObjectMixin's workspace_cache decorator."""
+
+    origin_type = "test-object-type"
+    _name = "test-object"
+
+    def __init__(self):
+        self.call_count = 0
+
+    @ObjectMixin.workspace_cache
+    def cached_method(self, *args, **kwargs):
+        self.call_count += 1
+        return self.call_count
+
+
+def test_workspace_cache():
+    obj = TestObject()
+    workspace = mock.Mock()
+    workspace.object_command_cache = {}
+
+    # First call
+    result1 = obj.cached_method(1, 2, key="value", workspace=workspace)
+    assert result1 == 1
+    assert obj.call_count == 1
+
+    # Second call with same arguments
+    result2 = obj.cached_method(1, 2, key="value", workspace=workspace)
+    assert result2 == 1
+    assert obj.call_count == 1
+
+    # Call with different args
+    result3 = obj.cached_method(3, 4, key="value", workspace=workspace)
+    assert result3 == 2
+    assert obj.call_count == 2
+
+    # Call with different kwargs
+    result4 = obj.cached_method(1, 2, key="new_value", workspace=workspace)
+    assert result4 == 3
+    assert obj.call_count == 3
+
+    # Verify cache contents
+    assert len(workspace.object_command_cache) == 3
+
+    # The caching assumes workspace= argument is present
+    with pytest.raises(KeyError, match="'workspace'"):
+        obj.cached_method(1, 2)

var/ramble/repos/builtin/package_managers/pip/package_manager.py

@@ -214,6 +214,10 @@ def get_spec_str(self, pkg, all_pkgs, compiler):
         """
         return pkg.spec
 
+    @PackageManagerBase.workspace_cache
+    def get_version(self, workspace=None):
+        return self.runner.get_version()
+
     def populate_inventory(
         self, workspace, force_compute=False, require_exist=False
     ):
@@ -223,7 +227,7 @@ def populate_inventory(
         self.runner.set_dry_run(workspace.dry_run)
         self.runner.configure_env(env_path)
 
-        pkgman_version = self.runner.get_version()
+        pkgman_version = self.get_version(workspace=workspace)
 
         self.app_inst.hash_inventory["package_manager"].append(
             {

var/ramble/repos/builtin/package_managers/spack-lightweight/package_manager.py

@@ -400,6 +400,10 @@ def _push_to_spack_cache(self, workspace, app_inst=None):
                 logger.die(e)
             pass
 
+    @PackageManagerBase.workspace_cache
+    def get_version(self, workspace=None):
+        return self.runner.get_version()
+
     def populate_inventory(
         self, workspace, force_compute=False, require_exist=False
     ):
@@ -411,7 +415,7 @@ def populate_inventory(
         self.runner.activate()
 
         try:
-            pkgman_version = self.runner.get_version()
+            pkgman_version = self.get_version(workspace=workspace)
         except RunnerError:
             pkgman_version = "unknown"