diff --git a/tmd/io/io.py b/tmd/io/io.py index b260679..d965f0d 100644 --- a/tmd/io/io.py +++ b/tmd/io/io.py @@ -122,12 +122,13 @@ def load_population(neurons, tree_types=None, name=None): pop = Population.Population(name=name) - files2load = [i for i in files if (i.endswith(".h5") or i.endswith(".swc"))] - - for i in files2load: + for filename in files: try: - pop.append_neuron(load_neuron(i, tree_types=tree_types)) + assert filename.endswith(('.h5', '.swc')) + pop.append_neuron(load_neuron(filename, tree_types=tree_types)) + except AssertionError: + raise Warning("{} is not a valid h5 or swc file".format(filename)) except LoadNeuronError: - print('File failed to load: {}'.format(i)) + print('File failed to load: {}'.format(filename)) return pop