Add "subclass match" for session.get_instance()

src/tribler/core/components/session.py

@@ -51,7 +51,20 @@ async def __aexit__(self, *_):
         await self.shutdown()
 
     def get_instance(self, comp_cls: Type[T]) -> Optional[T]:
-        return self.components.get(comp_cls)
+        # try to find a direct match
+        if direct_match := self.components.get(comp_cls):
+            return direct_match
+
+        # try to find a subclass match
+        candidates = {c for c in self.components if issubclass(c, comp_cls)}
+
+        if not candidates:
+            return None
+        if len(candidates) >= 2:
+            raise KeyError(f'Found more that two subclasses for the class {comp_cls}')
+
+        candidate = candidates.pop()
+        return self.components[candidate]
 
     def register(self, comp_cls: Type[Component], component: Component):
         if comp_cls in self.components:

src/tribler/core/components/tests/test_base_component.py

@@ -2,6 +2,7 @@
 
 from tribler.core.components.component import Component, MissedDependency, NoneComponent
 from tribler.core.components.session import Session
+from tribler.core.config.tribler_config import TriblerConfig
 
 
 class ComponentTestException(Exception):
@@ -20,16 +21,16 @@ async def run(self):
         async def shutdown(self):
             self.shutdown_was_executed = True
 
-    class ComponentA(TestComponent):
+    class TestComponentA(TestComponent):
         pass
 
-    class ComponentB(TestComponent):
+    class TestComponentB(TestComponent):
         pass
 
-    session = Session(tribler_config, [ComponentA(), ComponentB()])
+    session = Session(tribler_config, [TestComponentA(), TestComponentB()])
     async with session:
-        a = session.get_instance(ComponentA)
-        b = session.get_instance(ComponentB)
+        a = session.get_instance(TestComponentA)
+        b = session.get_instance(TestComponentB)
 
         for component in a, b:
             assert component.run_was_executed
@@ -44,18 +45,28 @@ class ComponentB(TestComponent):
         assert component.stopped
 
 
-async def test_required_dependency(tribler_config):
-    class ComponentA(Component):
-        pass
+class ComponentA(Component):
+    pass
 
-    class ComponentB(Component):
-        async def run(self):
-            await self.require_component(ComponentA)
 
-    session = Session(tribler_config, [ComponentA(), ComponentB()])
+class RequireA(Component):
+    async def run(self):
+        await self.require_component(ComponentA)
+
+
+class ComponentB(Component):
+    pass
+
+
+class DerivedB(ComponentB):
+    pass
+
+
+async def test_required_dependency(tribler_config):
+    session = Session(tribler_config, [ComponentA(), RequireA()])
     async with session:
         a = session.get_instance(ComponentA)
-        b = session.get_instance(ComponentB)
+        b = session.get_instance(RequireA)
 
         assert a in b.dependencies and not b.reverse_dependencies
         assert not a.dependencies and b in a.reverse_dependencies
@@ -67,49 +78,29 @@ async def run(self):
 
 
 async def test_required_dependency_missed(tribler_config):
-    class ComponentA(Component):
-        pass
-
-    class ComponentB(Component):
-        async def run(self):
-            await self.require_component(ComponentA)
-
-    session = Session(tribler_config, [ComponentB()])
-    with pytest.raises(MissedDependency, match='^Missed dependency: ComponentB requires ComponentA to be active$'):
+    session = Session(tribler_config, [RequireA()])
+    with pytest.raises(MissedDependency, match='^Missed dependency: RequireA requires ComponentA to be active$'):
         await session.start_components()
 
 
 async def test_required_dependency_missed_failfast(tribler_config):
-    class ComponentA(Component):
-        pass
-
-    class ComponentB(Component):
-        async def run(self):
-            await self.require_component(ComponentA)
-
-    session = Session(tribler_config, [ComponentB()], failfast=False)
+    session = Session(tribler_config, [RequireA()], failfast=False)
     async with session:
         await session.start_components()
-        b = session.get_instance(ComponentB)
+        b = session.get_instance(RequireA)
         assert b
         assert b.started_event.is_set()
         assert b.failed
 
 
 async def test_component_shutdown_failure(tribler_config):
-    class ComponentA(Component):
-        pass
-
-    class ComponentB(Component):
-        async def run(self):
-            await self.require_component(ComponentA)
-
+    class RequireAWithException(RequireA):
         async def shutdown(self):
             raise ComponentTestException
 
-    session = Session(tribler_config, [ComponentA(), ComponentB()])
+    session = Session(tribler_config, [ComponentA(), RequireAWithException()])
     a = session.get_instance(ComponentA)
-    b = session.get_instance(ComponentB)
+    b = session.get_instance(RequireAWithException)
 
     await session.start_components()
 
@@ -126,12 +117,6 @@ async def shutdown(self):
 
 
 async def test_maybe_component(loop, tribler_config):  # pylint: disable=unused-argument
-    class ComponentA(Component):
-        pass
-
-    class ComponentB(Component):
-        pass
-
     session = Session(tribler_config, [ComponentA()])
     async with session:
         component_a = await session.get_instance(ComponentA).maybe_component(ComponentA)
@@ -141,3 +126,27 @@ class ComponentB(Component):
         assert isinstance(component_b, NoneComponent)
         assert isinstance(component_b.any_attribute, NoneComponent)
         assert isinstance(component_b.any_attribute.any_nested_attribute, NoneComponent)
+
+
+def test_get_instance_direct_match(tribler_config: TriblerConfig):
+    session = Session(tribler_config, [ComponentA(), ComponentB(), DerivedB()])
+    assert isinstance(session.get_instance(ComponentB), ComponentB)
+
+
+def test_get_instance_subclass_match(tribler_config: TriblerConfig):
+    session = Session(tribler_config, [ComponentA(), DerivedB()])
+    assert isinstance(session.get_instance(ComponentB), DerivedB)
+
+
+def test_get_instance_no_match(tribler_config: TriblerConfig):
+    session = Session(tribler_config, [ComponentA()])
+    assert not session.get_instance(ComponentB)
+
+
+def test_get_instance_two_subclasses_match(tribler_config: TriblerConfig):
+    class SecondDerivedB(ComponentB):
+        pass
+
+    session = Session(tribler_config, [ComponentA(), DerivedB(), SecondDerivedB()])
+    with pytest.raises(KeyError):
+        session.get_instance(ComponentB)