#Python2 or Python3
#AspirinCode 2018
#Script that splits a multi-mol2 file into individual mol2 files.
#python split_multimol2.py multi-mol2.mol2 out_dir
import sys
import os
def split_multimol2(multimol2):
"""
Splits a multi-mol2 file.
Parameters
----------
multimol2 : str
Path to the multi-mol2 file.
Returns
----------
A generator object for lists for every extracted mol2-file. Lists contain
the molecule ID and the mol2 file contents.
e.g., ['ID1234', '@<TRIPOS>MOLECULE...'
"""
with open(multimol2, 'r') as mol2file:
line = mol2file.readline()
while not mol2file.tell() == os.fstat(mol2file.fileno()).st_size:
if line.startswith("@<TRIPOS>MOLECULE"):
mol2cont = []
mol2cont.append(line)
line = mol2file.readline()
molecule_id = line.strip()
while not line.startswith("@<TRIPOS>MOLECULE"):
mol2cont.append(line)
line = mol2file.readline()
if mol2file.tell() == os.fstat(mol2file.fileno()).st_size:
mol2cont.append(line)
break
mol2cont[-1] = mol2cont[-1].rstrip() # removes blank line at file end
yield [molecule_id, "".join(mol2cont)]
def write_multimol2(multimol2, out_dir):
"""
Splits a multi-mol2 file into smaller multi-mol2 files.
Parameters
-----------
multimol2 : str
Path to the multi-mol2 file.
out_dir : str:
Output directory. New files will be named
<molecule_name_1>.mol2, ... <molecule_name_n>.mol2
Returns
-----------
chunks : int
Number of files written.
"""
if not out_dir:
os.mkdir(out_dir)
single_mol2s = split_multimol2(args.MOL2_FILE)
for mol2 in single_mol2s:
out_mol2 = os.path.join(args.OUT_DIR, mol2[0]) + '.mol2'
with open(out_mol2, 'w') as out_file:
for line in mol2[1]:
out_file.write(line)
out_file.write('\n')
def write_multimol2_chunks(multimol2, chunk_size, out_dir):
"""
Splits a multi-mol2 file into smaller multi-mol2 files.
Parameters
-----------
multimol2 : str
Path to the multi-mol2 file.
chunksize : int
Number of mol2 files per chunk.
out_dir : str:
Output directory. New files will be named
<multimol2>_1.mol2, ... <multimol2>_n.mol2
Returns
-----------
chunks : int
Number of files written.
"""
if not os.path.exists(out_dir):
os.mkdir(out_dir)
out_path_stem = os.path.dirname(multimol2)
out_file_stem = os.path.basename(multimol2).split('.mol2')[0]
cnt = 0
chunks = 1
out_file = open(os.path.join(out_dir, out_file_stem)+'_%d.mol2' % chunks, 'w')
for mol2 in split_multimol2(multimol2):
cnt += 1
if cnt == chunk_size:
cnt = 0
chunks += 1
out_file.close()
out_file = open(os.path.join(out_dir, out_file_stem)+'_%d.mol2' % chunks, 'w')
out_file.write(mol2[1] + '\n')
out_file.close()
return chunks
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(
description='Splits a multi-mol2 file into individual mol2 files',
formatter_class=argparse.RawTextHelpFormatter
)
parser.add_argument('MOL2_FILE')
parser.add_argument('OUT_DIR')
parser.add_argument('-c', '--chunksize', help='Number of MOL2 structures per file (1 by default)', type=int)
parser.add_argument('-v', '--version', action='version', version='split_multimol2 v. 1.1')
args = parser.parse_args()
if args.chunksize:
write_multimol2_chunks(multimol2=args.MOL2_FILE, chunk_size=args.chunksize, out_dir=args.OUT_DIR)
else:
write_multimol2(multimol2=args.MOL2_FILE, out_dir=args.OUT_DIR)
