feat: introduce fairness

Author: triceo

timefold-solver-python-core/src/main/python/score/_constraint_stream.py

@@ -454,6 +454,20 @@ def concat(self, other):
         else:
             raise RuntimeError(f'Unhandled constraint stream type {type(other)}.')
 
+    def complement(self, item_type: Type[A]) -> 'UniConstraintStream[A]':
+        """
+        Adds to the stream all instances of a given class which are not yet present in it.
+        These instances must be present in the solution,
+        which means the class needs to be either a planning entity or a problem fact.
+
+        Parameters
+        ----------
+        item_type : Type[A]
+            the type of the instances to add to the stream.
+        """
+        item_type = get_class(item_type)
+        return UniConstraintStream(self.delegate.complement(item_type), self.package, self.a_type)
+
     def penalize(self, constraint_weight: ScoreType, match_weigher: Callable[[A], int] = None) -> \
             'UniConstraintBuilder[A, ScoreType]':
         """
@@ -1000,6 +1014,29 @@ def concat(self, other):
         else:
             raise RuntimeError(f'Unhandled constraint stream type {type(other)}.')
 
+    def complement(self, item_type: Type[A], padding_function: Callable[[A], B] = None) -> 'BiConstraintStream[A, B]':
+        """
+        Adds to the stream all instances of a given class which are not yet present in it.
+        These instances must be present in the solution,
+        which means the class needs to be either a planning entity or a problem fact.
+
+        The instances will be read from the first element of the input tuple.
+        When an output tuple needs to be created for the newly inserted instances,
+        the first element will be the new instance.
+        The rest of the tuple will be padded with the result of the padding function,
+        applied on the new instance.
+
+        Parameters
+        ----------
+        item_type : Type[A]
+            the type of the instances to add to the stream.
+
+        padding_function : Callable[[A], B]
+            a function that computes the padding value for the new tuple.
+        """
+        item_type = get_class(item_type)
+        return BiConstraintStream(self.delegate.complement(item_type, padding_function), self.package, self.a_type, self.b_type)
+
     def penalize(self, constraint_weight: ScoreType, match_weigher: Callable[[A, B], int] = None) -> \
             'BiConstraintBuilder[A, B, ScoreType]':
         """
@@ -1544,6 +1581,34 @@ def concat(self, other):
         else:
             raise RuntimeError(f'Unhandled constraint stream type {type(other)}.')
 
+    def complement(self, item_type: Type[A], padding_function_b: Callable[[A], B] = None,
+                   padding_function_c: Callable[[A], C] = None) -> 'TriConstraintStream[A, B, C]':
+        """
+        Adds to the stream all instances of a given class which are not yet present in it.
+        These instances must be present in the solution,
+        which means the class needs to be either a planning entity or a problem fact.
+
+        The instances will be read from the first element of the input tuple.
+        When an output tuple needs to be created for the newly inserted instances,
+        the first element will be the new instance.
+        The rest of the tuple will be padded with the result of the padding function,
+        applied on the new instance.
+
+        Parameters
+        ----------
+        item_type : Type[A]
+            the type of the instances to add to the stream.
+
+        padding_function_b : Callable[[A], B]
+            a function that computes the padding value for the second fact in the new tuple.
+
+        padding_function_c : Callable[[A], C]
+            a function that computes the padding value for the third fact in the new tuple.
+        """
+        item_type = get_class(item_type)
+        return TriConstraintStream(self.delegate.complement(item_type, padding_function_b, padding_function_c),
+                                   self.package, self.a_type, self.b_type, self.c_type)
+
     def penalize(self, constraint_weight: ScoreType,
                  match_weigher: Callable[[A, B, C], int] = None) -> 'TriConstraintBuilder[A, B, C, ScoreType]':
         """
@@ -2083,6 +2148,39 @@ def concat(self, other):
         else:
             raise RuntimeError(f'Unhandled constraint stream type {type(other)}.')
 
+    def complement(self, item_type: Type[A], padding_function_b: Callable[[A], B] = None,
+                   padding_function_c: Callable[[A], C] = None,
+                   padding_function_d: Callable[[A], D] = None) -> 'QuadConstraintStream[A, B, C, D]':
+        """
+        Adds to the stream all instances of a given class which are not yet present in it.
+        These instances must be present in the solution,
+        which means the class needs to be either a planning entity or a problem fact.
+
+        The instances will be read from the first element of the input tuple.
+        When an output tuple needs to be created for the newly inserted instances,
+        the first element will be the new instance.
+        The rest of the tuple will be padded with the result of the padding function,
+        applied on the new instance.
+
+        Parameters
+        ----------
+        item_type : Type[A]
+            the type of the instances to add to the stream.
+
+        padding_function_b : Callable[[A], B]
+            a function that computes the padding value for the second fact in the new tuple.
+
+        padding_function_c : Callable[[A], C]
+            a function that computes the padding value for the third fact in the new tuple.
+
+        padding_function_d : Callable[[A], D]
+            a function that computes the padding value for the fourth fact in the new tuple.
+        """
+        item_type = get_class(item_type)
+        return QuadConstraintStream(
+            self.delegate.complement(item_type, padding_function_b, padding_function_c, padding_function_d),
+            self.package, self.a_type, self.b_type, self.c_type, self.d_type)
+
     def penalize(self, constraint_weight: ScoreType,
                  match_weigher: Callable[[A, B, C, D], int] = None) -> 'QuadConstraintBuilder[A, B, C, D, ScoreType]':
         """

timefold-solver-python-core/src/main/python/score/_group_by.py

@@ -3,6 +3,7 @@
 from typing import Callable, Any, Sequence, TypeVar, List, Set, Dict, TYPE_CHECKING, overload
 if TYPE_CHECKING:
     from ai.timefold.solver.core.api.score.stream.common import SequenceChain
+    from ai.timefold.solver.core.api.score.stream.common import LoadBalance
     from ai.timefold.solver.core.api.score.stream.uni import UniConstraintCollector
     from ai.timefold.solver.core.api.score.stream.bi import BiConstraintCollector
     from ai.timefold.solver.core.api.score.stream.tri import TriConstraintCollector
@@ -135,13 +136,15 @@ class ConstraintCollectors:
     C = TypeVar('C')
     D = TypeVar('D')
     E = TypeVar('E')
+    Balanced = TypeVar('Balanced')
 
     #  Method return type variables
     A_ = TypeVar('A_')
     B_ = TypeVar('B_')
     C_ = TypeVar('C_')
     D_ = TypeVar('D_')
     E_ = TypeVar('E_')
+    Balanced_ = TypeVar('Balanced_')
 
     @staticmethod
     def _delegate():
@@ -993,6 +996,34 @@ def to_sorted_map(key_mapper, value_mapper, merge_function_or_set_creator=None):
         else:
             raise ValueError
 
+    @overload
+    @staticmethod
+    def load_balance(balanced_item_function: Callable[[A], Balanced], load_function: Callable[[A], int] = None,
+                     initial_load_function: Callable[[A], int] = None) -> \
+            'UniConstraintCollector[A, Any, LoadBalance[Balanced]]':
+        ...
+
+    @overload
+    @staticmethod
+    def load_balance(balanced_item_function: Callable[[A, B], Balanced], load_function: Callable[[A, B], int] = None,
+                     initial_load_function: Callable[[A, B], int] = None) -> \
+            'BiConstraintCollector[A, B, Any, LoadBalance[Balanced]]':
+        ...
+
+    @overload
+    @staticmethod
+    def load_balance(balanced_item_function: Callable[[A, B, C], Balanced], load_function: Callable[[A, B, C], int] = None,
+                     initial_load_function: Callable[[A, B, C], int] = None) -> \
+            'TriConstraintCollector[A, B, C, Any, LoadBalance[Balanced]]':
+        ...
+
+    @overload
+    @staticmethod
+    def load_balance(balanced_item_function: Callable[[A, B, C, D], Balanced],
+                     load_function: Callable[[A, B, C, D], int] = None,
+                     initial_load_function: Callable[[A, B, C, D], int] = None) -> \
+            'QuadConstraintCollector[A, B, C, D, Any, LoadBalance[Balanced]]':
+        ...
 
 # Must be at the bottom, constraint_stream depends on this module
 from ._constraint_stream import *