diff --git a/src/cosy_luigi/utils/traversals.py b/src/cosy_luigi/utils/traversals.py index 6e279ff..1b323c5 100644 --- a/src/cosy_luigi/utils/traversals.py +++ b/src/cosy_luigi/utils/traversals.py @@ -4,13 +4,13 @@ import inspect from abc import ABC -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING, cast from cosy_luigi import CoSyLuigiTask if TYPE_CHECKING: - from collections.abc import Iterable, Sequence + from collections.abc import Sequence def flatten( @@ -32,8 +32,12 @@ def flatten( for task_or_task_collection in heterogeneous_task_collection for task in ( flatten(*cast("Iterable[type[CoSyLuigiTask]]", task_or_task_collection)) - if isinstance(task_or_task_collection, (tuple, list)) - else cast("type[CoSyLuigiTask]", task_or_task_collection).get_all_variants() + if isinstance(task_or_task_collection, Iterable) + else ( + variant + for variant in cast("type[CoSyLuigiTask]", task_or_task_collection).get_all_variants() + if not (inspect.isabstract(variant) or ABC in variant.__bases__) + ) if inspect.isabstract(task_or_task_collection) or ABC in cast("type[CoSyLuigiTask]", task_or_task_collection).__bases__ else (task_or_task_collection,) diff --git a/tests/test_abstract_variant_expansion.py b/tests/test_abstract_variant_expansion.py index 5d07d2b..4a50315 100644 --- a/tests/test_abstract_variant_expansion.py +++ b/tests/test_abstract_variant_expansion.py @@ -17,6 +17,14 @@ class ABCInheritedTask(CoSyLuigiTask, ABC): """An abstract class. This class is abstract because it inherits from ABC.""" +class ABCTaskFromABCInherited(ABCInheritedTask, ABC): + """An abstract class that implements ABCInheritedTask.""" + + +class DeeperConcreteTaskFromABCTaskFromABCInherited(ABCTaskFromABCInherited): + """A class that indirectly implements ABCInheritedTask by inheriting from ABCTaskFromABCInherited.""" + + class ConcreteTaskFromABCInherited(ABCInheritedTask): """A class that implements ABCInheritedTask.""" @@ -81,7 +89,11 @@ def test_expansion_from_abc(): """Tests if adding a class that is abstract because it inherits from ABC expands to all of its subclasses when added to a CoSyLuigiRepo.""" repo = CoSyLuigiRepo(ABCInheritedTask) - assert repo.luigi_repo == {ConcreteTaskFromABCInherited, DeeperConcreteTaskFromABCInherited} + assert repo.luigi_repo == { + ConcreteTaskFromABCInherited, + DeeperConcreteTaskFromABCInherited, + DeeperConcreteTaskFromABCTaskFromABCInherited, + } def test_expansion_from_abstract(): @@ -99,6 +111,7 @@ def test_expansion_from_abc_and_abstract(): repo = CoSyLuigiRepo(ABCInheritedTask, AbstractTask) assert repo.luigi_repo == { ConcreteTaskFromABCInherited, + DeeperConcreteTaskFromABCTaskFromABCInherited, DeeperConcreteTaskFromABCInherited, ConcreteTaskFromAbstract, DeeperConcreteTaskFromAbstract,