From 03769be12677b575dd9fc87026cf4fde2724c27e Mon Sep 17 00:00:00 2001 From: Tias Guns Date: Sat, 4 Apr 2026 22:36:51 +0200 Subject: [PATCH 1/2] args typing for Minimum, with helpers --- cpmpy/expressions/globalfunctions.py | 27 +++++++++++------ cpmpy/expressions/utils.py | 44 ++++++++++++++++++++++++++-- 2 files changed, 59 insertions(+), 12 deletions(-) diff --git a/cpmpy/expressions/globalfunctions.py b/cpmpy/expressions/globalfunctions.py index b45063bdb..26e8b0a0a 100644 --- a/cpmpy/expressions/globalfunctions.py +++ b/cpmpy/expressions/globalfunctions.py @@ -79,8 +79,8 @@ def decompose(self): from ..exceptions import CPMpyException, IncompleteFunctionError, TypeError from .core import Expression, Operator, ExprLike, ListLike -from .variables import boolvar, intvar, cpm_array -from .utils import flatlist, argval, is_num, is_int, eval_comparison, is_any_list, is_boolexpr, get_bounds, argvals, implies +from .variables import intvar, NDVarArray +from .utils import flatlist, argval, is_num, eval_comparison, is_any_list, is_boolexpr, get_bounds, argvals, implies, argvals_listlike, get_bounds_listlike class GlobalFunction(Expression): @@ -162,18 +162,27 @@ def __init__(self, arg_list: ListLike[ExprLike]): Arguments: arg_list (ListLike[ExprLike]): List of expressions or constants of which to compute the minimum """ - super().__init__("min", tuple(flatlist(arg_list))) + if isinstance(arg_list, NDVarArray): + has_subexpr = arg_list.has_subexpr() + return super().__init__("min", tuple(arg_list.flat), has_subexpr=has_subexpr) + + # arg: tuple[ExprLike, ...] + super().__init__("min", tuple(arg_list)) + + @property + def args(self) -> tuple[ExprLike, ...]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args def value(self) -> Optional[int]: """ Returns: Optional[int]: The minimum value of the arguments, or None if any argument is not assigned """ - vargs = [argval(a) for a in self.args] - if any(val is None for val in vargs): + vals = argvals_listlike(self.args) + if vals is None: return None - - return min(vargs) + return min(vals) def decompose(self) -> tuple[Expression, list[Expression]]: """ @@ -196,8 +205,8 @@ def get_bounds(self) -> tuple[int, int]: Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the minimum value """ - bnds = [get_bounds(x) for x in self.args] - return min(lb for lb, ub in bnds), min(ub for lb, ub in bnds) + lbs, ubs = get_bounds_listlike(self.args) + return min(lbs), min(ubs) class Maximum(GlobalFunction): diff --git a/cpmpy/expressions/utils.py b/cpmpy/expressions/utils.py index 6a4890a1b..35c41d642 100644 --- a/cpmpy/expressions/utils.py +++ b/cpmpy/expressions/utils.py @@ -35,7 +35,7 @@ import math from collections.abc import Iterable # for flatten from itertools import combinations -from typing import TYPE_CHECKING, TypeGuard, Union, Optional +from typing import TYPE_CHECKING, TypeGuard, Optional from cpmpy.exceptions import IncompleteFunctionError if TYPE_CHECKING: @@ -152,6 +152,25 @@ def argvals(arr): return [argvals(arg) for arg in arr] return argval(arr) +def argvals_listlike(lst: ListLike[ExprLike]) -> Optional[list[int]]: + """ The well-typed way to get the values of a list of ExprLike's, or None if any expression is not assigned """ + _Expr = cp.expressions.core.Expression + + vals: list[int] = [] + for e in lst: + if isinstance(e, _Expr): + v = e.value() + if v is None: + return None + vals.append(v) + elif isinstance(e, int): + vals.append(e) + else: # only np.integer is still in ExprLike + vals.append(int(e)) + return vals + + + def eval_comparison(str_op, lhs, rhs): """ @@ -197,7 +216,7 @@ def get_bounds(expr): # from cpmpy.expressions.core import Expression # from cpmpy.expressions.variables import cpm_array - if isinstance(expr, (cp.expressions.core.Expression, cp.variables.NDVarArray)): + if isinstance(expr, (cp.expressions.core.Expression, cp.expressions.variables.NDVarArray)): return expr.get_bounds() elif is_any_list(expr): lbs, ubs = zip(*[get_bounds(e) for e in expr]) @@ -207,10 +226,29 @@ def get_bounds(expr): if is_bool(expr): return int(expr), int(expr) return math.floor(expr), math.ceil(expr) + +def get_bounds_listlike(lst: ListLike[ExprLike]) -> tuple[list[int], list[int]]: + """ The well-typed way to get the bounds of a list of ExprLike's """ + _Expr = cp.expressions.core.Expression + + lbs: list[int] = [] + ubs: list[int] = [] + for e in lst: + if isinstance(e, _Expr): + (lb, ub) = e.get_bounds() + lbs.append(lb) + ubs.append(ub) + elif isinstance(e, int): + lbs.append(e) + ubs.append(e) + else: # only np.integer is still in ExprLike + lbs.append(int(e)) + ubs.append(int(e)) + return lbs, ubs def implies(expr, other): """ like :func:`~cpmpy.expressions.core.Expression.implies`, but also safe to use for non-expressions """ - if isinstance(expr, (cp.expressions.core.Expression, cp.variables.NDVarArray)): + if isinstance(expr, (cp.expressions.core.Expression, cp.expressions.variables.NDVarArray)): # both implement .implies() return expr.implies(other) elif is_true_cst(expr): From 939e99de0f7fdeedb102e106b2e708a14962df25 Mon Sep 17 00:00:00 2001 From: Tias Guns Date: Tue, 7 Apr 2026 13:01:04 +0200 Subject: [PATCH 2/2] and rest of global functions (all but Multiplication and Element) --- cpmpy/expressions/core.py | 2 + cpmpy/expressions/globalfunctions.py | 262 +++++++++++++++++---------- 2 files changed, 172 insertions(+), 92 deletions(-) diff --git a/cpmpy/expressions/core.py b/cpmpy/expressions/core.py index fdda69d52..fd1942fba 100644 --- a/cpmpy/expressions/core.py +++ b/cpmpy/expressions/core.py @@ -413,6 +413,8 @@ def __pow__(self, other: Any, modulo: Optional[int] = None): raise TypeError("Power operator: modulo not supported") if not isinstance(other, (int, np.integer)): raise TypeError(f"Power operator requires a constant integer exponent, not: {other}") + if other == 0: + return 1 if other == 1: return self return cp.Power(self, other) diff --git a/cpmpy/expressions/globalfunctions.py b/cpmpy/expressions/globalfunctions.py index 26e8b0a0a..176bfbed5 100644 --- a/cpmpy/expressions/globalfunctions.py +++ b/cpmpy/expressions/globalfunctions.py @@ -73,14 +73,14 @@ def decompose(self): """ import warnings # for deprecation warning -from typing import Optional +from typing import Any, Iterable, Optional, cast import numpy as np import cpmpy as cp from ..exceptions import CPMpyException, IncompleteFunctionError, TypeError -from .core import Expression, Operator, ExprLike, ListLike +from .core import Expression, Operator, Comparison, ExprLike, ListLike from .variables import intvar, NDVarArray -from .utils import flatlist, argval, is_num, eval_comparison, is_any_list, is_boolexpr, get_bounds, argvals, implies, argvals_listlike, get_bounds_listlike +from .utils import argval, is_num, eval_comparison, is_any_list, is_boolexpr, get_bounds, argvals, implies, argvals_listlike, get_bounds_listlike class GlobalFunction(Expression): @@ -219,18 +219,28 @@ def __init__(self, arg_list: ListLike[ExprLike]): Arguments: arg_list (ListLike[ExprLike]): List of expressions or constants of which to compute the maximum """ - super().__init__("max", tuple(flatlist(arg_list))) + if isinstance(arg_list, NDVarArray): + has_subexpr = arg_list.has_subexpr() + return super().__init__("max", tuple(arg_list.flat), has_subexpr=has_subexpr) + + # arg: tuple[ExprLike, ...] + super().__init__("max", tuple(arg_list)) + + @property + def args(self) -> tuple[ExprLike, ...]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args def value(self) -> Optional[int]: """ Returns: Optional[int]: The maximum value of the arguments, or None if any argument is not assigned """ - vargs = [argval(a) for a in self.args] - if any(val is None for val in vargs): + vals = argvals_listlike(self.args) + if vals is None: return None - return max(vargs) + return max(vals) def decompose(self) -> tuple[Expression, list[Expression]]: """ @@ -253,8 +263,8 @@ def get_bounds(self) -> tuple[int, int]: Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the maximum value """ - bnds = [get_bounds(x) for x in self.args] - return max(lb for lb, ub in bnds), max(ub for lb, ub in bnds) + lbs, ubs = get_bounds_listlike(self.args) + return max(lbs), max(ubs) class Abs(GlobalFunction): @@ -267,18 +277,24 @@ def __init__(self, expr: Expression): Arguments: expr (Expression): Expression of which to compute the absolute value """ + # args: tuple[Expression] super().__init__("abs", (expr,)) + + @property + def args(self) -> tuple[Expression]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args def value(self) -> Optional[int]: """ Returns: Optional[int]: The absolute value of the argument, or None if the argument is not assigned """ - varg = argval(self.args[0]) - if varg is None: + val = self.args[0].value() + if val is None: return None - return abs(varg) + return abs(val) def decompose(self) -> tuple[Expression, list[Expression]]: """ @@ -292,13 +308,13 @@ def decompose(self) -> tuple[Expression, list[Expression]]: tuple[Expression, list[Expression]]: A tuple containing the expression representing the absolute value (may be the argument itself, its negation, or an auxiliary variable), and a list of constraints defining it (empty if no auxiliary variable is needed) """ arg = self.args[0] - lb, ub = get_bounds(arg) + lb, ub = arg.get_bounds() if lb >= 0: # always positive return arg, [] if ub <= 0: # always negative return -arg, [] - _abs = intvar(*self.get_bounds()) + _abs = intvar(*self.get_bounds()) # not just lb,ub, see implementation return _abs, [(arg >= 0).implies(_abs == arg), (arg < 0).implies(_abs == -arg)] def get_bounds(self) -> tuple[int, int]: @@ -308,7 +324,7 @@ def get_bounds(self) -> tuple[int, int]: Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the absolute value """ - lb, ub = get_bounds(self.args[0]) + lb, ub = self.args[0].get_bounds() if lb >= 0: return lb, ub if ub <= 0: @@ -344,7 +360,7 @@ def __init__(self, x: ExprLike, y: ExprLike): super().__init__("mul", (x, y)) self.is_lhs_num = is_lhs_num - def update_args(self, args): + def update_args(self, args, has_subexpr: Optional[bool] = None): """ Allows in-place update of the expression's arguments. Resets all cached computations which depend on the expression tree. """ @@ -356,7 +372,7 @@ def update_args(self, args): (x, y) = (y, x) is_lhs_num = True - super().update_args((x, y)) + super().update_args((x, y), has_subexpr) self.is_lhs_num = is_lhs_num def __repr__(self): @@ -459,9 +475,15 @@ def __init__(self, x: ExprLike, y: ExprLike): x (ExprLike): Expression or constant to divide y (ExprLike): Expression or constant to divide by """ + # args: tuple[ExprLike, ExprLike] super().__init__("div", (x, y)) + + @property + def args(self) -> tuple[ExprLike, ExprLike]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args - def __repr__(self): + def __repr__(self) -> str: """ Returns: str: String representation of integer division as 'x div y' @@ -470,7 +492,7 @@ def __repr__(self): return "{} div {}".format(f"({x})" if isinstance(x, Expression) else x, f"({y})" if isinstance(y, Expression) else y) - def decompose(self): + def decompose(self) -> tuple[Expression, list[Expression]]: """ Decomposition of Integer Division global function, rounding towards zero. @@ -481,24 +503,31 @@ def decompose(self): tuple[Expression, list[Expression]]: A tuple containing the auxiliary variable representing the integer division, and a list of constraints defining it """ x,y = self.args - safen = [] + defining: list[Expression] = [] - y_lb, y_ub = get_bounds(y) - if y_lb <= 0 <= y_ub: - safen = [y != 0] - warnings.warn(f"Division constraint is unsafe, and will be forced to be total by this decomposition. If you are using {self} in a nested context, this is not valid, and you need to safen first using cpmpy.transformations.safening.no_partial_functions") + if isinstance(y, Expression): + y_lb, y_ub = y.get_bounds() + if y_lb <= 0 <= y_ub: + defining.append(y != 0) + warnings.warn(f"Division constraint is unsafe, and will be forced to be total by this decomposition. If you are using {self} in a nested context, this is not valid, and you need to safen first using cpmpy.transformations.safening.no_partial_functions") r = intvar(*get_bounds(x % y)) # remainder _div = intvar(*self.get_bounds()) - return _div, safen + [(x == (y * _div) + r), abs(r) < abs(y), abs(y) * abs(_div) <= abs(x)] + defining.append(Comparison("==", x, (y * _div) + r)) + defining.extend([abs(r) < abs(y), abs(y) * abs(_div) <= abs(x)]) + return _div, defining - def value(self): + def value(self) -> Optional[int]: """ Returns: int: The integer division of the arguments, or None if the arguments are not assigned """ - x,y = argvals(self.args) - if x is None or y is None: + x,y = self.args + x = argval(x) + if x is None: + return None + y = argval(y) + if y is None: return None try: @@ -508,7 +537,7 @@ def value(self): + "\n Use argval(expr) to get the value of expr with relational " "semantics.") - def get_bounds(self): + def get_bounds(self) -> tuple[int, int]: """ Returns the bounds of the Division global function @@ -554,9 +583,15 @@ def __init__(self, x: ExprLike, y: ExprLike): x (ExprLike): Expression or constant for the dividend y (ExprLike): Expression or constant for the divisor """ + # args: tuple[ExprLike, ExprLike] super().__init__("mod", (x, y)) + + @property + def args(self) -> tuple[ExprLike, ExprLike]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args - def __repr__(self): + def __repr__(self) -> str: """ Returns: str: String representation with 'mod' as notation @@ -565,7 +600,7 @@ def __repr__(self): return "{} mod {}".format(f"({x})" if isinstance(x, Expression) else x, f"({y})" if isinstance(y, Expression) else y) - def decompose(self): + def decompose(self) -> tuple[Expression, list[Expression]]: """ Decomposition of Modulo global function, using integer division (rounding towards zero) @@ -574,28 +609,34 @@ def decompose(self): https://marcelkliemannel.com/articles/2021/dont-confuse-integer-division-with-floor-division/ """ x,y = self.args - safen = [] + defining: list[Expression] = [] - y_lb, y_ub = get_bounds(y) - if y_lb <= 0 <= y_ub: - safen = [y != 0] + if isinstance(y, Expression): + y_lb, y_ub = y.get_bounds() + if y_lb <= 0 <= y_ub: + defining.append(y != 0) warnings.warn(f"Modulo constraint is unsafe, and will be forced to be total by this decomposition. If you are using {self} in a nested context, this is not valid, and you need to safen first using cpmpy.transformations.safening.no_partial_functions") _mod = intvar(*self.get_bounds()) k = intvar(*get_bounds((x - _mod) // y)) # integer quotient (multiplier) - return _mod, safen + [ + defining.extend([ k * y + _mod == x, # module is remainder of integer division abs(_mod) < abs(y), # remainder is smaller than divisor x * _mod >= 0 # remainder is negative iff x is negative - ] + ]) + return _mod, defining - def value(self): + def value(self) -> Optional[int]: """ Returns: int: The modulo of the arguments, or None if the arguments are not assigned """ - x,y = argvals(self.args) - if x is None or y is None: + x,y = self.args + x = argval(x) + if x is None: + return None + y = argval(y) + if y is None: return None try: # modulo defined with integer division @@ -605,24 +646,25 @@ def value(self): + "\n Use argval(expr) to get the value of expr with relational " "semantics.") - def get_bounds(self): + def get_bounds(self) -> tuple[int, int]: """ Returns the bounds of the Modulo global function Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the modulo """ - lb1, ub1 = get_bounds(self.args[0]) - lb2, ub2 = get_bounds(self.args[1]) - if lb2 == ub2 == 0: - raise ZeroDivisionError("Domain of {} only contains 0".format(self.args[1])) + x,y = self.args + x_lb, x_ub = get_bounds(x) + y_lb, y_ub = get_bounds(y) + if y_lb == y_ub == 0: + raise ZeroDivisionError(f"Domain of {y} only contains 0") # the (abs of) the maximum value of the remainder is always one smaller than the absolute value of the divisor - lb = lb2 + (lb2 <= 0) - (lb2 >= 0) - ub = ub2 + (ub2 <= 0) - (ub2 >= 0) - if lb1 >= 0: # result will be positive if first argument is positive + lb = y_lb + (y_lb <= 0) - (y_lb >= 0) + ub = y_ub + (y_ub <= 0) - (y_ub >= 0) + if x_lb >= 0: # result will be positive if first argument is positive return 0, max(-lb, ub, 0) # lb = 0 - elif ub1 <= 0: # result will be negative if first argument is negative + elif x_ub <= 0: # result will be negative if first argument is negative return min(-ub, lb, 0), 0 # ub = 0 return min(-ub, lb, 0), max(-lb, ub, 0) # 0 should always be in the domain @@ -634,19 +676,27 @@ class Power(GlobalFunction): Only non-negative constant integer exponents are supported. """ - def __init__(self, base: ExprLike, exponent: int|np.integer): + def __init__(self, base: Expression, exponent: int|np.integer): """ Arguments: base (ExprLike): Expression or constant to raise to the power exponent (int | np.integer): Non-negative integer exponent (constant only, no variable) """ - if not is_num(exponent): - raise TypeError(f"Power constraint takes an integer number as second argument, not: {exponent}") - if exponent < 0: - raise ValueError(f"Power constraint only supports non-negative integer exponents, not: {exponent}") + if not isinstance(exponent, int): + exponent = int(exponent) + #raise TypeError(f"Power constraint takes an integer number as second argument, not: {exponent}") + if exponent <= 0: + raise ValueError(f"Power constraint only supports positive integer exponents, not: {exponent}") + + # args: tuple[Expression, int] super().__init__("pow", (base, exponent)) + + @property + def args(self) -> tuple[Expression, int]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args - def decompose(self): + def decompose(self) -> tuple[Expression, list[Expression]]: """ Decomposition of Power global function, using integer multiplication. @@ -656,26 +706,25 @@ def decompose(self): tuple[Expression, list[Expression]]: A tuple containing the auxiliary variable representing the power, and a list of constraints defining it """ base, exp = self.args - if exp == 0: - return 1,[] _pow = base for _ in range(1,exp): _pow *= base return _pow,[] - def value(self): + def value(self) -> Optional[int]: """ Returns: int: The power of the arguments, or None if the arguments are not assigned """ - base, exp = argvals(self.args) - if base is None or exp is None: + base, exp = self.args + base_val = base.value() + if base_val is None: return None - return base**exp + return base_val**exp - def get_bounds(self): + def get_bounds(self) -> tuple[int, int]: """ Returns the bounds of the Power global function @@ -683,7 +732,7 @@ def get_bounds(self): tuple[int, int]: A tuple of (lower bound, upper bound) for the power """ base, exp = self.args - lb_base, ub_base = get_bounds(base) + lb_base, ub_base = base.get_bounds() bounds = [lb_base ** exp, ub_base ** exp] return min(bounds), max(bounds) @@ -829,12 +878,16 @@ def __init__(self, arr: ListLike[ExprLike], val: ExprLike): arr (ListLike[ExprLike]): List of expressions or constants to count in val (ExprLike): 'Value' to count occurences of (can also be an expression) """ - if not is_any_list(arr): - raise TypeError(f"Count(arr, val) takes an array of expressions as first argument, not: {arr}") - if is_any_list(val): + if isinstance(val, (list, np.ndarray)): raise TypeError(f"Count(arr, val) takes a numeric expression as second argument, not a list: {val}") + # args: tuple[ListLike[ExprLike], ExprLike] super().__init__("count", (arr, val)) + @property + def args(self) -> tuple[ListLike[ExprLike], ExprLike]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args + def decompose(self) -> tuple[Expression, list[Expression]]: """ Decomposition of the Count global function. @@ -857,8 +910,8 @@ def value(self) -> Optional[int]: if vval is None: return None - varr = [argval(a) for a in arr] - if any(v is None for v in varr): + varr = argvals_listlike(arr) + if varr is None: return None return sum((a == vval) for a in varr) @@ -870,7 +923,7 @@ def get_bounds(self) -> tuple[int, int]: Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the count value """ - arr, val = self.args + arr = self.args[0] return 0, len(arr) @@ -890,11 +943,19 @@ def __init__(self, arr: ListLike[ExprLike], vals: ListLike[int|np.integer]): arr (ListLike[ExprLike]): List of expressions or constants to count occurrences in vals (ListLike[int | np.integer]): List of integer constants whose occurrences are counted """ - if not is_any_list(arr) or not is_any_list(vals): - raise TypeError(f"Among takes as input two arrays, not: {arr} and {vals}") - if any(isinstance(val, Expression) for val in vals): - raise TypeError(f"Among takes a set of integer values as input, not {vals}") - super().__init__("among", (arr, vals)) + has_subexpr = None + if isinstance(arr, NDVarArray): # quick check + has_subexpr = arr.has_subexpr() + + vals = [int(v) for v in vals] # only pure Python ints + + # args: tuple[ListLike[ExprLike], list[int]] + super().__init__("among", (arr, vals), has_subexpr=has_subexpr) + + @property + def args(self) -> tuple[ListLike[ExprLike], list[int]]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args def decompose(self) -> tuple[Expression, list[Expression]]: """ @@ -915,8 +976,8 @@ def value(self) -> Optional[int]: Optional[int]: The number of variables in arr that take a value present in vals, or None if any element in arr is not assigned """ arr, vals = self.args - varr = argvals(arr) # recursive handling of nested structures - if any(v is None for v in varr): + varr = argvals_listlike(arr) # recursive handling of nested structures + if varr is None: return None return int(sum(np.isin(varr, vals))) @@ -928,7 +989,7 @@ def get_bounds(self) -> tuple[int, int]: Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the among count value """ - arr, vals = self.args + arr = self.args[0] return 0, len(arr) @@ -945,9 +1006,17 @@ def __init__(self, arr: ListLike[ExprLike]): Arguments: arr (ListLike[ExprLike]): List of expressions or constants to count distinct values in """ - if not is_any_list(arr): - raise ValueError(f"NValue(arr) takes an array as input, not: {arr}") - super().__init__("nvalue", tuple(arr)) + has_subexpr = None + if isinstance(arr, NDVarArray): # quick check + has_subexpr = arr.has_subexpr() + + # args: tuple[ExprLike, ...] + super().__init__("nvalue", tuple(arr), has_subexpr=has_subexpr) + + @property + def args(self) -> tuple[ExprLike, ...]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args def decompose(self) -> tuple[Expression, list[Expression]]: """ @@ -966,7 +1035,7 @@ def decompose(self) -> tuple[Expression, list[Expression]]: Returns: tuple[Expression, list[Expression]]: A tuple containing the sum expression representing the number of distinct values, and an empty list of constraints (no auxiliary variables needed) """ - lbs, ubs = get_bounds(self.args) + lbs, ubs = get_bounds_listlike(self.args) lb, ub = min(lbs), max(ubs) return cp.sum(cp.any(a == v for a in self.args) for v in range(lb, ub+1)), [] @@ -976,8 +1045,8 @@ def value(self) -> Optional[int]: Returns: Optional[int]: The number of distinct values in the array, or None if any element in arr is not assigned """ - vargs = [argval(a) for a in self.args] - if any(v is None for v in vargs): + vargs = argvals_listlike(self.args) + if vargs is None: return None return len(set(vargs)) @@ -1008,11 +1077,20 @@ def __init__(self, arr: ListLike[ExprLike], n: int|np.integer): arr (ListLike[ExprLike]): List of expressions or constants to count distinct values in n (int | np.integer): Integer constant to exclude from the count """ - if not is_any_list(arr): - raise ValueError("NValueExcept takes an array as input") - if not is_num(n): - raise ValueError(f"NValueExcept takes an integer as second argument, but got {n} of type {type(n)}") - super().__init__("nvalue_except", (arr, n)) + has_subexpr = None + if isinstance(arr, NDVarArray): # quick check + has_subexpr = arr.has_subexpr() + + if not isinstance(n, int): + n = int(n) + + # args: tuple[ListLike[ExprLike], int] + super().__init__("nvalue_except", (arr, n), has_subexpr=has_subexpr) + + @property + def args(self) -> tuple[ListLike[ExprLike], int]: + """ READ-ONLY, well-typed argument of this global function """ + return self._args def decompose(self) -> tuple[Expression, list[Expression]]: """ @@ -1027,7 +1105,7 @@ def decompose(self) -> tuple[Expression, list[Expression]]: """ arr, n = self.args - lbs, ubs = get_bounds(arr) + lbs, ubs = get_bounds_listlike(arr) lb, ub = min(lbs), max(ubs) return cp.sum([cp.any(a == v for a in arr) for v in range(lb, ub+1) if v != n]), [] @@ -1038,8 +1116,8 @@ def value(self) -> Optional[int]: Optional[int]: The number of distinct values in the array, excluding value n, or None if any element in arr is not assigned """ arr, n = self.args - varr = [argval(a) for a in arr] - if any(v is None for v in varr): + varr = argvals_listlike(arr) + if varr is None: return None if n in varr: @@ -1054,5 +1132,5 @@ def get_bounds(self) -> tuple[int, int]: Returns: tuple[int, int]: A tuple of (lower bound, upper bound) for the number of distinct values (excluding n) """ - arr, n = self.args + arr = self.args[0] return 0, len(arr)