Another way to handle multiple Inputs cache policies (#16786)

Co-authored-by: nate nowack <[email protected]>

Author: desertaxle

src/prefect/cache_policies.py

@@ -157,39 +157,15 @@ class CompoundCachePolicy(CachePolicy):
 
     policies: list[CachePolicy] = field(default_factory=list)
 
-    def __sub__(self, other: str) -> CachePolicy:
-        if not isinstance(other, str):  # type: ignore[reportUnnecessaryIsInstance]
-            raise TypeError("Can only subtract strings from key policies.")
-
-        # Subtract from each sub-policy; this may stack new Inputs(...) layers
-        new_subpolicies: list[CachePolicy] = [p - other for p in self.policies]
-
-        # If there are multiple Inputs, unify them all into one
-        inputs_policies = [p for p in new_subpolicies if isinstance(p, Inputs)]
+    def __post_init__(self) -> None:
+        # deduplicate any Inputs policies
+        inputs_policies = [p for p in self.policies if isinstance(p, Inputs)]
+        self.policies = [p for p in self.policies if not isinstance(p, Inputs)]
         if inputs_policies:
-            # Gather all excludes into a single set
             all_excludes: set[str] = set()
             for inputs_policy in inputs_policies:
                 all_excludes.update(inputs_policy.exclude)
-
-            # Remove all old Inputs from the subpolicy list
-            new_subpolicies = [p for p in new_subpolicies if not isinstance(p, Inputs)]
-
-            # Append one merged Inputs that excludes the union
-            merged_inputs = Inputs(exclude=sorted(all_excludes))
-            new_subpolicies.append(merged_inputs)
-
-        # If no sub‐policies exist after this, create an Inputs with [other]
-        # (to preserve the "always add an Inputs" behavior)
-        if not new_subpolicies:
-            new_subpolicies = [Inputs(exclude=[other])]
-
-        return CompoundCachePolicy(
-            policies=new_subpolicies,
-            key_storage=self.key_storage,
-            isolation_level=self.isolation_level,
-            lock_manager=self.lock_manager,
-        )
+            self.policies.append(Inputs(exclude=sorted(all_excludes)))
 
     def compute_key(
         self,
@@ -212,6 +188,28 @@ def compute_key(
             return None
         return hash_objects(*keys, raise_on_failure=True)
 
+    def __add__(self, other: "CachePolicy") -> "CachePolicy":
+        # Call the superclass add method to handle validation
+        super().__add__(other)
+
+        if isinstance(other, CompoundCachePolicy):
+            policies = [*self.policies, *other.policies]
+        else:
+            policies = [*self.policies, other]
+
+        return CompoundCachePolicy(
+            policies=policies,
+            key_storage=self.key_storage or other.key_storage,
+            isolation_level=self.isolation_level or other.isolation_level,
+            lock_manager=self.lock_manager or other.lock_manager,
+        )
+
+    def __sub__(self, other: str) -> "CachePolicy":
+        if not isinstance(other, str):  # type: ignore[reportUnnecessaryIsInstance]
+            raise TypeError("Can only subtract strings from key policies.")
+        new = Inputs(exclude=[other])
+        return CompoundCachePolicy(policies=[*self.policies, new])
+
 
 @dataclass
 class _None(CachePolicy):

tests/test_cache_policies.py

@@ -158,6 +158,16 @@ def test_nones_are_ignored(self):
         )
         assert compound_key is None
 
+    def test_adding_two_compound_policies_merges_policies(self):
+        one = CompoundCachePolicy(policies=[Inputs(), TaskSource()])
+        two = CompoundCachePolicy(policies=[RunId()])
+        policy = one + two
+        assert isinstance(policy, CompoundCachePolicy)
+        assert len(policy.policies) == 3
+        assert Inputs() in policy.policies
+        assert RunId() in policy.policies
+        assert TaskSource() in policy.policies
+
     def test_compound_policy_deduplicates_inputs_on_subtraction(self):
         """Regression test for https://github.com/PrefectHQ/prefect/issues/16773"""
         # Create a compound policy with multiple Inputs policies
@@ -169,6 +179,8 @@ def test_compound_policy_deduplicates_inputs_on_subtraction(self):
                 Inputs(exclude=["y"]),
             ]
         )
+        # Inputs get combined into a single policy
+        assert len(policy.policies) == 2
 
         # Subtract a new key
         new_policy = policy - "z"
@@ -184,8 +196,7 @@ def test_compound_policy_deduplicates_inputs_on_subtraction(self):
         # So we should have one merged Inputs policy and one CompoundCachePolicy containing TaskSource
         assert len(new_policy.policies) == 2
         assert any(
-            isinstance(p, CompoundCachePolicy)
-            and any(isinstance(sp, TaskSource) for sp in p.policies)
+            isinstance(p, CompoundCachePolicy) or isinstance(p, TaskSource)
             for p in new_policy.policies
         )