diff --git a/modelling/pymod/_raw_model.py b/modelling/pymod/_raw_model.py index b08ae45b3c51846a0303069c4646c1e5c64ae2fa..50bb31defea8058bd06acfaa603e6ad653d92d0b 100644 --- a/modelling/pymod/_raw_model.py +++ b/modelling/pymod/_raw_model.py @@ -77,8 +77,11 @@ def BuildRawModel(aln, chain_names = None, include_ligands = False, If *aln* is of type :class:`ost.seq.AlignmentHandle`, *chain_names* is expected to be a :class:`str`. If *aln* is of type :class:`ost.seq.AlignmentList`, - chain_names is expected to be a :class:`list` of - :class:`str` of same size as *aln*. + *chain_names* is expected to be a :class:`list` of + :class:`str` of same size as *aln* or a :class:`str`. + For the latter case, chains will consecutively named + according to characters in *chain_names*. + :type chain_names: :class:`str` / :class:`list` :param spdbv_style: True, if we need a model in the old SPDBV style. @@ -130,9 +133,15 @@ def BuildRawModel(aln, chain_names = None, include_ligands = False, raise RuntimeError('Number of alns and chain_names must be '\ 'consistent') name_list = chain_names + elif isinstance(chain_names, str): + if len(chain_names) < len(aln_list): + raise RuntimeError('If you provide a string as chain_names, '\ + 'it must be at least as long as the '\ + 'AlignmentList in aln.') + name_list = [n for n in chain_names[:len(aln_list)]] else: - raise RuntimeError('chain_names must be list if aln is of type '\ - 'ost.seq.AlignmentList') + raise RuntimeError('chain_names must be list of str or str if aln '\ + 'is of type ost.seq.AlignmentList') else: raise RuntimeError('aln must be of type ost.seq.AlignmentHandle or '\ 'ost.seq.AlignmentList') diff --git a/modelling/tests/test_modelling.py b/modelling/tests/test_modelling.py index f8e8b946b4120d24b820c041d12a6a1d7b65ec3d..3e99059ac00f92af1691bc0ed406d306d9fda6a7 100644 --- a/modelling/tests/test_modelling.py +++ b/modelling/tests/test_modelling.py @@ -76,24 +76,28 @@ class ModellingTests(unittest.TestCase): result = modelling.BuildRawModel(aln_lst, ['cheese', 'steak']) self.assertEqual(result.model.chains[0].GetName(), 'cheese') self.assertEqual(result.model.chains[1].GetName(), 'steak') + result = modelling.BuildRawModel(aln_lst, 'ch') + self.assertEqual(result.model.chains[0].GetName(), 'c') + self.assertEqual(result.model.chains[1].GetName(), 'h') + result = modelling.BuildRawModel(aln_lst, 'cheese') + self.assertEqual(result.model.chains[0].GetName(), 'c') + self.assertEqual(result.model.chains[1].GetName(), 'h') # we only accept a string as chain_names if aln is AlignmentHandle self.assertRaises(RuntimeError, modelling.BuildRawModel, aln, 1) self.assertRaises(RuntimeError, modelling.BuildRawModel, aln, ['A']) - # we only accept a list as chain_names if aln is AlignmentList + # we only accept a list or str as chain_names if aln is AlignmentList self.assertRaises(RuntimeError, modelling.BuildRawModel, aln_lst, 1) - self.assertRaises(RuntimeError, modelling.BuildRawModel, aln_lst, 'A') # size also matters... self.assertRaises(RuntimeError, modelling.BuildRawModel, aln_lst, ['A']) + self.assertRaises(RuntimeError, modelling.BuildRawModel, aln_lst, 'A') # increase size of aln_list => at some point we should run out of # default chain_names aln_lst = 100*[aln] self.assertRaises(RuntimeError, modelling.BuildRawModel, aln_lst) - - def testModeledSequence(self): # test if the model has the sequence we want. tpl = io.LoadPDB('data/gly.pdb')