diff --git a/python/cuopt/cuopt/linear_programming/problem.py b/python/cuopt/cuopt/linear_programming/problem.py index 49885acf52..68f4fbe99b 100644 --- a/python/cuopt/cuopt/linear_programming/problem.py +++ b/python/cuopt/cuopt/linear_programming/problem.py @@ -35,6 +35,27 @@ class VType(str, Enum): SEMI_CONTINUOUS = VType.SEMI_CONTINUOUS +def _to_vtype(value): + """ + Coerces a variable type to a :py:class:`VType` member. + Besides VType members, the single character codes are accepted as ``str`` + or ``bytes``. + """ + if isinstance(value, VType): + return value + try: + # UnicodeDecodeError and the enum lookup failure are both ValueError. + return VType( + value.decode() if isinstance(value, (bytes, bytearray)) else value + ) + except ValueError: + valid = ", ".join(repr(t.value) for t in VType) + raise ValueError( + f"Invalid variable type {value!r}. Expected a VType member or " + f"one of {valid}." + ) from None + + class CType(str, Enum): """ The sense of a constraint is either LE, GE or EQ. @@ -94,8 +115,10 @@ class Variable: ---------- VariableName : str Name of the Variable. - VariableType : CONTINUOUS, INTEGER, or SEMI_CONTINUOUS - Variable type. + VariableType : VType + Variable type, always normalized to a :py:class:`VType` member + (CONTINUOUS, INTEGER, or SEMI_CONTINUOUS). Assigning a ``str`` or + ``bytes`` character code converts it; anything else raises ValueError. LB : float Lower Bound of the Variable. UB : float @@ -129,6 +152,14 @@ def __init__( self.VariableName = vname self.MIPStart = float("nan") + @property + def VariableType(self): + return self._variable_type + + @VariableType.setter + def VariableType(self, value): + self._variable_type = _to_vtype(value) + def getIndex(self): """ Get the index position of the variable in the problem. @@ -181,13 +212,15 @@ def getUpperBound(self): def setVariableType(self, val): """ Sets the variable type of the variable. - Variable types can be CONTINUOUS, INTEGER, or SEMI_CONTINUOUS. + Variable types can be CONTINUOUS, INTEGER, or SEMI_CONTINUOUS, or the + equivalent character code as ``str`` or ``bytes``. + Raises ValueError for any other value. """ self.VariableType = val def getVariableType(self): """ - Returns the type of the variable. + Returns the type of the variable as a :py:class:`VType` member. """ return self.VariableType @@ -2078,7 +2111,7 @@ def NumNZs(self): def IsMIP(self): # Returns if the problem is a MIP problem. for var in self.vars: - if var.VariableType in ("I", "S", b"I", b"S"): + if var.VariableType in (INTEGER, SEMI_CONTINUOUS): return True return False diff --git a/python/cuopt/cuopt/tests/linear_programming/test_python_API.py b/python/cuopt/cuopt/tests/linear_programming/test_python_API.py index 964809553f..9780cb6e88 100644 --- a/python/cuopt/cuopt/tests/linear_programming/test_python_API.py +++ b/python/cuopt/cuopt/tests/linear_programming/test_python_API.py @@ -160,6 +160,45 @@ def test_constraint_duplicate_terms_slack(): assert c.compute_slack() == pytest.approx(6.0) +def test_variable_type_is_normalized(): + prob = Problem() + from_enum = prob.addVariable(vtype=INTEGER) + from_str = prob.addVariable(vtype="I") + from_bytes = prob.addVariable(vtype=b"I") + default = prob.addVariable() + + for var in (from_enum, from_str, from_bytes): + assert var.VariableType is VType.INTEGER + assert default.VariableType is VType.CONTINUOUS + assert prob.IsMIP + + # Both the setter and direct assignment normalize. + from_str.setVariableType(b"S") + assert from_str.VariableType is VType.SEMI_CONTINUOUS + from_bytes.VariableType = "C" + assert from_bytes.VariableType is VType.CONTINUOUS + + with pytest.raises(ValueError): + from_enum.setVariableType(7) + + +def test_variable_type_normalized_from_mps(tmp_path): + prob = Problem("mip") + x = prob.addVariable(lb=0.0, ub=10.0, vtype=INTEGER, name="x") + y = prob.addVariable(lb=0.0, ub=10.0, name="y") + prob.addConstraint(x + y <= 5, name="c") + prob.setObjective(x + y, sense=MAXIMIZE) + + path = str(tmp_path / "mip.mps") + prob.writeMPS(path) + + loaded = Problem.read(path) + types = [v.VariableType for v in loaded.getVariables()] + assert all(isinstance(t, VType) for t in types) + assert VType.INTEGER in types + assert loaded.IsMIP + + def test_semi_continuous_variable(): prob = Problem("Semi-continuous") x = prob.addVariable(lb=5.0, ub=10.0, vtype=SEMI_CONTINUOUS, name="x")