parser.add_argument('-o', '--output', type=str, help='MOL2 output file for filtered results.') parser.add_argument('-c', '--criteria', type=str, help='Query atom and charge ranges e.g., "O.2,-1.2,2;O.3,-20.0,100.0".') parser.add_argument('-m', '--matchall', action='store_true', help='If flag is provided, molecule must satisfy all criteria.') args = parser.parse_args() if not args.input: print('Please provide a valid input file. Run ./mol2_filter_funcgroups.py -h for help.') quit() if not args.output: print('Please provide a output file path. Run ./mol2_filter_funcgroups.py -h for help.') quit() if not args.criteria: print('Please provide some search criteria. Run ./mol2_filter_funcgroups.py -h for help.') quit() clist = create_chargetype_list(group_charge_pairs=args.criteria, atom_list=None) mmol2 = split_multimol2(multimol2=args.input) filterfunc = count_matches if args.matchall: filterfunc = match_all with open(args.output, 'w') as out_file: for m in mmol2: res = filterfunc(mol2_cont=m[1].split('\n'), chargetype_list=clist) if res: for line in m[1]: out_file.write(line) print(m[0])
def test_create_chargetype_list(): l1 = create_chargetype_list(group_charge_pairs='O.2,-1.2,2;O.3,-20.0,100.0' ,atom_list=None) assert l1 == [['O.2', -1.2, 2.0], ['O.3', -20.0, 100.0]]