Add support for anonymous enums and enum constructor functions

Author: Dekker1

CHANGELOG.rst

@@ -14,6 +14,8 @@ Added
 
 - Add additional ``Driver.executable`` property to safely access the location
   of the executable used by the ``Driver`` object.
+- Support for input and output of complex MiniZinc enumerated types, such as
+  anonymous enumerated types or enumerated types using constructor functions.
 
 Removed
 ^^^^^^^

src/minizinc/__init__.py

@@ -12,6 +12,7 @@
 from .model import Method, Model
 from .result import Result, Status
 from .solver import Solver
+from .types import AnonEnum, ConstrEnum
 
 logger = logging.getLogger("minizinc")
 
@@ -36,6 +37,8 @@
 __all__ = [
     "__version__",
     "default_driver",
+    "AnonEnum",
+    "ConstrEnum",
     "Driver",
     "Instance",
     "Method",

src/minizinc/json.py

@@ -8,6 +8,7 @@
 from json import JSONDecoder, JSONEncoder, loads
 
 from .error import MiniZincWarning, error_from_stream_obj
+from .types import AnonEnum, ConstrEnum
 
 try:
     import numpy
@@ -19,6 +20,10 @@ class MZNJSONEncoder(JSONEncoder):
     def default(self, o):
         if isinstance(o, Enum):
             return {"e": o.name}
+        if isinstance(o, AnonEnum):
+            return {"e": o.enumName, "i": o.value}
+        if isinstance(o, ConstrEnum):
+            return {"c": o.constructor, "e": o.argument}
         if isinstance(o, set) or isinstance(o, range):
             return {"set": [{"e": i.name} if isinstance(i, Enum) else i for i in o]}
         if numpy is not None:
@@ -37,9 +42,20 @@ def __init__(self, enum_map=None, *args, **kwargs):
             self.enum_map = enum_map
         JSONDecoder.__init__(self, object_hook=self.mzn_object_hook, *args, **kwargs)
 
+    def transform_enum_object(self, obj):
+        # TODO: This probably is an enum, but could still be a record
+        if "e" in obj:
+            if len(obj) == 1:
+                return self.enum_map.get(obj["e"], obj["e"])
+            elif len(obj) == 2 and "c" in obj:
+                return ConstrEnum(obj["c"], obj["e"])
+            elif len(obj) == 2 and "i" in obj:
+                return AnonEnum(obj["e"], obj["i"])
+        return obj
+
     def mzn_object_hook(self, obj):
-        if isinstance(obj, dict) and len(obj) == 1:
-            if "set" in obj:
+        if isinstance(obj, dict):
+            if len(obj) == 1 and "set" in obj:
                 if len(obj["set"]) == 1 and isinstance(obj["set"][0], list):
                     assert len(obj["set"][0]) == 2
                     return range(obj["set"][0][0], obj["set"][0][1] + 1)
@@ -49,13 +65,13 @@ def mzn_object_hook(self, obj):
                     if isinstance(item, list):
                         assert len(item) == 2
                         li.extend([i for i in range(item[0], item[1] + 1)])
-                    elif isinstance(item, dict) and len(item) == 1 and "e" in item:
-                        li.append(self.enum_map.get(item["e"], item["e"]))
+                    elif isinstance(item, dict):
+                        li.append(self.transform_enum_object(item))
                     else:
                         li.append(item)
                 return set(li)
-            elif "e" in obj:
-                return self.enum_map.get(obj["e"], obj["e"])
+            else:
+                return self.transform_enum_object(obj)
         return obj
 
 

src/minizinc/types.py

@@ -0,0 +1,28 @@
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class AnonEnum:
+    """Representation of anonymous enumeration values in MiniZinc"""
+
+    enumName: str
+    value: int
+
+    def __str__(self):
+        return f"to_enum({self.enumName},{self.value})"
+
+
+@dataclass(frozen=True)
+class ConstrEnum:
+    """Representation of constructor function enumerated values in MiniZinc"""
+
+    constructor: str
+    argument: Any
+
+    def __str__(self):
+        return f"{self.constructor}({self.argument})"

tests/test_types.py

@@ -8,6 +8,7 @@
 
 from minizinc import Instance
 from minizinc.result import Status
+from minizinc.types import AnonEnum, ConstrEnum
 
 
 class TestEnum(InstanceTestCase):
@@ -104,6 +105,36 @@ def test_intenum_collections(self):
         # TODO: assert result["arr_t"] == [TT(3), TT(2), TT(1)]
         assert result["set_t"] == {TT(1), TT(2)}
 
+    def test_constructor_enum(self):
+        self.instance = Instance(self.solver)
+        self.instance.add_string(
+            """
+            enum T = X(1..3);
+            var T: x;
+            constraint x > X(1) /\\ x < X(3); % TODO: Remove for MiniZinc 2.7+
+            """
+        )
+        # TODO: Remove for MiniZinc 2.7+
+        # self.instance["x"] = ConstrEnum("X", 2)
+        result = self.instance.solve()
+        assert isinstance(result["x"], ConstrEnum)
+        assert result["x"] == ConstrEnum("X", 2)
+        assert str(result["x"]) == "X(2)"
+
+    def test_anon_enum(self):
+        self.instance = Instance(self.solver)
+        self.instance.add_string(
+            """
+            enum T = _(1..5);
+            var T: x;
+            """
+        )
+        self.instance["x"] = AnonEnum("T", 3)
+        result = self.instance.solve()
+        assert isinstance(result["x"], AnonEnum)
+        assert result["x"].value == 3
+        assert str(result["x"]) == "to_enum(T,3)"
+
 
 class TestSets(InstanceTestCase):
     def test_ranges(self):