8
8
from json import JSONDecoder , JSONEncoder , loads
9
9
10
10
from .error import MiniZincWarning , error_from_stream_obj
11
+ from .types import AnonEnum , ConstrEnum
11
12
12
13
try :
13
14
import numpy
@@ -19,6 +20,10 @@ class MZNJSONEncoder(JSONEncoder):
19
20
def default (self , o ):
20
21
if isinstance (o , Enum ):
21
22
return {"e" : o .name }
23
+ if isinstance (o , AnonEnum ):
24
+ return {"e" : o .enumName , "i" : o .value }
25
+ if isinstance (o , ConstrEnum ):
26
+ return {"c" : o .constructor , "e" : o .argument }
22
27
if isinstance (o , set ) or isinstance (o , range ):
23
28
return {"set" : [{"e" : i .name } if isinstance (i , Enum ) else i for i in o ]}
24
29
if numpy is not None :
@@ -37,9 +42,20 @@ def __init__(self, enum_map=None, *args, **kwargs):
37
42
self .enum_map = enum_map
38
43
JSONDecoder .__init__ (self , object_hook = self .mzn_object_hook , * args , ** kwargs )
39
44
45
+ def transform_enum_object (self , obj ):
46
+ # TODO: This probably is an enum, but could still be a record
47
+ if "e" in obj :
48
+ if len (obj ) == 1 :
49
+ return self .enum_map .get (obj ["e" ], obj ["e" ])
50
+ elif len (obj ) == 2 and "c" in obj :
51
+ return ConstrEnum (obj ["c" ], obj ["e" ])
52
+ elif len (obj ) == 2 and "i" in obj :
53
+ return AnonEnum (obj ["e" ], obj ["i" ])
54
+ return obj
55
+
40
56
def mzn_object_hook (self , obj ):
41
- if isinstance (obj , dict ) and len ( obj ) == 1 :
42
- if "set" in obj :
57
+ if isinstance (obj , dict ):
58
+ if len ( obj ) == 1 and "set" in obj :
43
59
if len (obj ["set" ]) == 1 and isinstance (obj ["set" ][0 ], list ):
44
60
assert len (obj ["set" ][0 ]) == 2
45
61
return range (obj ["set" ][0 ][0 ], obj ["set" ][0 ][1 ] + 1 )
@@ -49,13 +65,13 @@ def mzn_object_hook(self, obj):
49
65
if isinstance (item , list ):
50
66
assert len (item ) == 2
51
67
li .extend ([i for i in range (item [0 ], item [1 ] + 1 )])
52
- elif isinstance (item , dict ) and len ( item ) == 1 and "e" in item :
53
- li .append (self .enum_map . get (item [ "e" ], item [ "e" ] ))
68
+ elif isinstance (item , dict ):
69
+ li .append (self .transform_enum_object (item ))
54
70
else :
55
71
li .append (item )
56
72
return set (li )
57
- elif "e" in obj :
58
- return self .enum_map . get (obj [ "e" ], obj [ "e" ] )
73
+ else :
74
+ return self .transform_enum_object (obj )
59
75
return obj
60
76
61
77
0 commit comments