diff --git a/tangos/input_handlers/finding.py b/tangos/input_handlers/finding.py index bafff942..8973fb45 100644 --- a/tangos/input_handlers/finding.py +++ b/tangos/input_handlers/finding.py @@ -6,6 +6,7 @@ from .. import config from ..log import logger +from ..util.subclasses import find_subclasses def find(extension=None, mtd=None, ignore=None, basename="", patterns=[]): @@ -46,16 +47,11 @@ def best_matching_handler(cls, basename): handler_names = [] handler_timestep_lengths = [] base = os.path.join(config.base, basename) - if len(cls.__subclasses__()) == 0: - return cls # Add all subclasses, sub-subclasses, sub-subclasses, ... - all_possible_handlers = [] - subclasses = cls.__subclasses__() - while subclasses: - subclass = subclasses.pop() - subclasses.extend(subclass.__subclasses__()) - all_possible_handlers.append(subclass) + all_possible_handlers = find_subclasses(cls) + if not all_possible_handlers: + return cls for possible_handler in all_possible_handlers: timesteps_detected = find(basename=base + "/", patterns=possible_handler.patterns) diff --git a/tangos/util/subclasses.py b/tangos/util/subclasses.py new file mode 100644 index 00000000..c0ddc424 --- /dev/null +++ b/tangos/util/subclasses.py @@ -0,0 +1,11 @@ +from typing import List + +def find_subclasses(cls: type) -> List[type]: + subclasses = cls.__subclasses__() + all_subclasses = [] + while subclasses: + subclass = subclasses.pop() + subclasses.extend(subclass.__subclasses__()) + all_subclasses.append(subclass) + + return all_subclasses \ No newline at end of file