33
44def test_serialization ():
55 from ortools .sat .python import cp_model
6+ from ortools .sat import cp_model_pb2
67 from google .protobuf import text_format
78 from pathlib import Path
89
@@ -13,28 +14,55 @@ def _detect_binary_mode(filename: str) -> bool:
1314 return True
1415 raise ValueError (f"Unknown extension for file: { filename } " )
1516
17+ # Changed in ortools 9.15: was model.Proto().SerializeToString() / text_format.MessageToString()
18+ # Now use model.export_to_file() which auto-detects format by extension
1619 def export_model (
1720 model : cp_model .CpModel , filename : str , binary : bool | None = None
1821 ):
1922 binary = _detect_binary_mode (filename ) if binary is None else binary
20- if binary :
21- Path (filename ).write_bytes (model .Proto ().SerializeToString ())
23+ # export_to_file uses .txt extension for text format, otherwise binary
24+ # So we need to handle the mismatch for some extensions
25+ if binary and filename .endswith (".txt" ):
26+ # Force binary even with .txt extension - use temp file
27+ temp_file = filename + ".pb"
28+ model .export_to_file (temp_file )
29+ Path (filename ).write_bytes (Path (temp_file ).read_bytes ())
30+ Path (temp_file ).unlink ()
31+ elif not binary and not filename .endswith (".txt" ):
32+ # Force text even without .txt extension
33+ temp_file = filename + ".txt"
34+ model .export_to_file (temp_file )
35+ Path (filename ).write_text (Path (temp_file ).read_text ())
36+ Path (temp_file ).unlink ()
2237 else :
23- Path ( filename ). write_text ( text_format . MessageToString ( model .Proto ()) )
38+ model .export_to_file ( filename )
2439
40+ # Changed in ortools 9.15: was model.Proto().ParseFromString() / text_format.Parse()
41+ # Now use model.Proto().parse_text_format() for text, or cp_model_pb2 for binary
2542 def import_model (filename : str , binary : bool | None = None ) -> cp_model .CpModel :
2643 binary = _detect_binary_mode (filename ) if binary is None else binary
2744 model = cp_model .CpModel ()
2845 if binary :
29- model .Proto ().ParseFromString (Path (filename ).read_bytes ())
46+ # Parse binary via standard protobuf, then convert to text for import
47+ proto = cp_model_pb2 .CpModelProto ()
48+ proto .ParseFromString (Path (filename ).read_bytes ())
49+ model .Proto ().parse_text_format (text_format .MessageToString (proto ))
3050 else :
31- text_format . Parse ( Path (filename ).read_text (), model . Proto ())
51+ model . Proto (). parse_text_format ( Path (filename ).read_text ())
3252 return model
3353
3454 def rename_variable_names (model : cp_model .CpModel ):
35- for i , var in enumerate (model .proto .variables ):
55+ for i , var in enumerate (model .Proto () .variables ):
3656 var .name = f"x{ i } "
3757
58+ # Changed in ortools 9.15: Proto() objects no longer support == comparison
59+ # Compare via text representation instead
60+ def protos_equal (m1 : cp_model .CpModel , m2 : cp_model .CpModel ) -> bool :
61+ with tempfile .TemporaryDirectory () as td :
62+ m1 .export_to_file (f"{ td } /m1.txt" )
63+ m2 .export_to_file (f"{ td } /m2.txt" )
64+ return Path (f"{ td } /m1.txt" ).read_text () == Path (f"{ td } /m2.txt" ).read_text ()
65+
3866 model = cp_model .CpModel ()
3967 x = [model .NewIntVar (0 , 10 , f"x{ i } " ) for i in range (10 )]
4068 model .add (sum (x ) <= 20 )
@@ -43,16 +71,16 @@ def rename_variable_names(model: cp_model.CpModel):
4371 with tempfile .TemporaryDirectory () as tmpdir :
4472 export_model (model , f"{ tmpdir } /model.pb" )
4573 model2 = import_model (f"{ tmpdir } /model.pb" )
46- assert model . Proto () == model2 . Proto ( )
74+ assert protos_equal ( model , model2 )
4775 export_model (model , f"{ tmpdir } /model.txt" , binary = False )
4876 model3 = import_model (f"{ tmpdir } /model.txt" , binary = False )
49- assert model . Proto () == model3 . Proto ( )
77+ assert protos_equal ( model , model3 )
5078 export_model (model , f"{ tmpdir } /model.pbtxt" , binary = False )
5179 model4 = import_model (f"{ tmpdir } /model.pbtxt" , binary = False )
52- assert model . Proto () == model4 . Proto ( )
80+ assert protos_equal ( model , model4 )
5381 export_model (model , f"{ tmpdir } /model.pb.txt" , binary = True )
5482 model5 = import_model (f"{ tmpdir } /model.pb.txt" , binary = True )
55- assert model . Proto () == model5 . Proto ( )
83+ assert protos_equal ( model , model5 )
5684 rename_variable_names (model )
5785 export_model (model , f"{ tmpdir } /model_renamed.pb.txt" )
5886 model6 = import_model (f"{ tmpdir } /model_renamed.pb.txt" )
0 commit comments