def write_pb2(cmpd, filename, binary=True):
"""Convert mb.Compound to Protobuf Message file.
Parameters
----------
cmpd : mb.Compound
filename : str
binary: bool, default True
If True, will print a binary file
If False, will print to a text file
"""
cmpd_to_proto = {}
root_proto = compound_pb2.Compound()
root_proto = _mb_to_proto(cmpd, root_proto)
cmpd_to_proto[cmpd] = root_proto
for sub_cmpd in cmpd.successors():
parent_cmpd = sub_cmpd.parent
sub_proto = cmpd_to_proto[parent_cmpd].children.add()
sub_proto = _mb_to_proto(sub_cmpd, sub_proto)
cmpd_to_proto[sub_cmpd] = sub_proto
_add_proto_bonds(cmpd, root_proto)
if binary:
with open(filename, "wb") as f:
f.write(root_proto.SerializeToString())
else:
with open(filename, "w") as f:
PrintMessage(root_proto, f)
def read_pb2(filename, binary=True):
""" Convert a Protobuf Message file into mb.Compound
Parameters
---------
filename : str
binary: bool, default True
If True, will print a binary file
If False, will print to a text file
Todo: This could be more elegantly detected
Returns
------
root_compound : mb.Compound
"""
root_proto = compound_pb2.Compound()
if binary:
with open(filename, 'rb') as f:
root_proto.ParseFromString(f.read())
else:
with open(filename, 'r') as f:
Merge(f.read(), root_proto)
proto_to_cmpd = {}
root_compound = _proto_to_mb(root_proto)
proto_to_cmpd[root_proto.id] = root_compound
for sub_proto, parent_proto in _proto_successors(root_proto):
if parent_proto.id not in proto_to_cmpd:
parent_cmpd = _proto_to_mb(parent_proto)
proto_to_cmpd[parent_proto.id] = parent_cmpd
parent_cmpd = proto_to_cmpd[parent_proto.id]
if sub_proto.id not in proto_to_cmpd:
sub_cmpd = _proto_to_mb(sub_proto)
proto_to_cmpd[sub_proto.id] = sub_cmpd
sub_cmpd = proto_to_cmpd[sub_proto.id]
parent_cmpd.add(sub_cmpd)
_add_mb_bonds(root_proto, root_compound, proto_to_cmpd)
return root_compound
def write_pb2(cmpd, filename, binary=True):
""" Convert mb.Compound to Protobuf Message file
Parameters
---------
cmpd : mb.Compound
filename : str
binary: bool, default True
If True, will print a binary file
If False, will print to a text file
Todo: This could be more elegantly detected
Notes
----
Todo: Handle Ports in the protocol buffer (.proto) and in this writer/reader
"""
cmpd_to_proto = {}
root_proto = compound_pb2.Compound()
root_proto = _mb_to_proto(cmpd, root_proto)
cmpd_to_proto[cmpd] = root_proto
for sub_cmpd in cmpd.successors():
parent_cmpd = sub_cmpd.parent
sub_proto = cmpd_to_proto[parent_cmpd].children.add()
sub_proto = _mb_to_proto(sub_cmpd, sub_proto)
cmpd_to_proto[sub_cmpd] = sub_proto
_add_proto_bonds(cmpd, root_proto)
if binary:
with open(filename, 'wb') as f:
f.write(root_proto.SerializeToString())
else:
with open(filename, 'w') as f:
PrintMessage(root_proto, f)
