update variable type

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
 

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")