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)