cascache.py: Add properties to CaptureTree requests

1 files changed, 9 insertions, 2 deletions

diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index 9c0b46d1e..519de6e2e 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -26,6 +26,7 @@ import ctypes
 import multiprocessing
 import signal
 import time
+from typing import Optional, List
 
 import grpc
 
@@ -34,7 +35,7 @@ from .._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from .._protos.build.buildgrid import local_cas_pb2
 
 from .. import _signals, utils
-from ..types import FastEnum
+from ..types import FastEnum, SourceRef
 from .._exceptions import CASCacheError
 
 from .casdprocessmanager import CASDProcessManager
@@ -339,15 +340,21 @@ class CASCache:
     #
     # Args:
     #     path (str): Path to directory to import
+    #     properties Optional[List[str]]: List of properties to request
     #
     # Returns:
     #     (Digest): The digest of the imported directory
     #
-    def import_directory(self, path):
+    def import_directory(self, path: str, properties: Optional[List[str]] = None) -> SourceRef:
         local_cas = self.get_local_cas()
 
         request = local_cas_pb2.CaptureTreeRequest()
         request.path.append(path)
+
+        if properties:
+            for _property in properties:
+                request.node_properties.append(_property)
+
         response = local_cas.CaptureTree(request)
 
         if len(response.responses) != 1: