diff --git a/barrOs_library.py b/barrOs_library.py index 6ae2b0e4a0def457cfe9943f0c2fa8301cda2925..deda67759480a4688e92dea8e404cc71724c0abc 100644 --- a/barrOs_library.py +++ b/barrOs_library.py @@ -20,7 +20,7 @@ from random import * from scipy import stats ## DEFINE DEFAULT PARAMETERS -accepted_input_types = ['pdb_file', 'pdbID', 'hhpred', 'pdbsum_fasta'] +accepted_input_types = ['pdb_file', 'pdbID', 'pdb_folder', 'hhpred', 'pdbsum_fasta'] accepted_modes = ['membrane', 'non-membrane', 'all'] ## 0. FOR LOGGING @@ -333,6 +333,9 @@ def get_all_pdbIDs(input_files, input_types): if input_types[i] == 'pdb_file': pdb_file = input_string pdbIDs.append(pdb_file) + elif input_types[i] == 'pdb_folder': + pdb_files = ['{}/{}'.format(input_string, f) for f in os.listdir(input_string) if f.endswith('.pdb')] + pdbIDs += pdb_files elif input_types[i] == 'pdbID': pdbIDs.append(input_string) elif input_types[i] == 'hhpred':