fix(foundation): let register_extractor_* accept pcapkit's own built-in classes

pcapkit/foundation/extraction.py

@@ -26,15 +26,15 @@
 from pcapkit.corekit.io import SeekableReader
 from pcapkit.corekit.module import ModuleDescriptor
 from pcapkit.dumpkit.common import make_dumper
-from pcapkit.foundation.engines.engine import Engine
+from pcapkit.foundation.engines.engine import Engine, EngineBase
 from pcapkit.foundation.engines.pcap import PCAP as PCAP_Engine
 from pcapkit.foundation.engines.pcapng import PCAPNG as PCAPNG_Engine
 from pcapkit.foundation.reassembly import ReassemblyManager
 from pcapkit.foundation.reassembly.data import ReassemblyData
-from pcapkit.foundation.reassembly.reassembly import Reassembly
+from pcapkit.foundation.reassembly.reassembly import Reassembly, ReassemblyBase
 from pcapkit.foundation.traceflow import TraceFlowManager
 from pcapkit.foundation.traceflow.data import TraceFlowData
-from pcapkit.foundation.traceflow.traceflow import TraceFlow
+from pcapkit.foundation.traceflow.traceflow import TraceFlow, TraceFlowBase
 from pcapkit.utilities.exceptions import (CallableError, FileNotFound, FormatError, IterableError,
                                           RegistryError, UnsupportedCall, stacklevel)
 from pcapkit.utilities.logging import get_logger
@@ -387,7 +387,12 @@ def register_engine(cls, name: 'str', engine: 'ModuleDescriptor[Engine] | Type[E
         """
         if isinstance(engine, ModuleDescriptor):
             engine = engine.klass
-        if not issubclass(engine, Engine):
+        # NOTE: checked against ``EngineBase`` rather than ``Engine``: every built-in
+        # engine subclasses the base directly (``engines/pcap.py`` imports it as
+        # ``EngineBase as Engine``) precisely so that it is *not* auto-registered by
+        # ``Engine.__init_subclass__``, which made this check reject pcapkit's own
+        # classes. ``Engine`` is itself an ``EngineBase``, so this only widens. See #513.
+        if not issubclass(engine, EngineBase):
             raise RegistryError(f'engine must be an Engine subclass, not {engine!r}')
         if name in cls.__engine__:
             warn(f'engine {name} already registered, overwriting', RegistryWarning)
@@ -409,7 +414,9 @@ def register_reassembly(cls, protocol: 'str', reassembly: 'ModuleDescriptor[Reas
         """
         if isinstance(reassembly, ModuleDescriptor):
             reassembly = reassembly.klass
-        if not issubclass(reassembly, Reassembly):
+        # NOTE: ``ReassemblyBase`` rather than ``Reassembly``, for the reason given in
+        # :meth:`register_engine` above -- see #513.
+        if not issubclass(reassembly, ReassemblyBase):
             raise RegistryError(f'reassembly must be a Reassembly subclass, not {reassembly!r}')
         if protocol in cls.__reassembly__:
             warn(f'reassembly {protocol} already registered, overwriting', RegistryWarning)
@@ -431,7 +438,9 @@ def register_traceflow(cls, protocol: 'str', traceflow: 'ModuleDescriptor[TraceF
         """
         if isinstance(traceflow, ModuleDescriptor):
             traceflow = traceflow.klass
-        if not issubclass(traceflow, TraceFlow):
+        # NOTE: ``TraceFlowBase`` rather than ``TraceFlow``, for the reason given in
+        # :meth:`register_engine` above -- see #513.
+        if not issubclass(traceflow, TraceFlowBase):
             raise RegistryError(f'traceflow must be a TraceFlow subclass, not {traceflow!r}')
         if protocol in cls.__traceflow__:
             warn(f'traceflow {protocol} already registered, overwriting', RegistryWarning)

tests/foundation/registry/test_foundation.py

@@ -103,6 +103,74 @@ def test_callback_and_extractor_registration_wrappers(self) -> None:
                                                   'TCP_TraceFlow')
         self.assertIsInstance(register.call_args.args[1], registry.ModuleDescriptor)
 
+    def test_registration_accepts_pcapkit_own_builtin_classes(self) -> None:
+        """The built-ins pass the ``issubclass`` gate, and non-subclasses still fail.
+
+        This is the regression for GitHub issue #513. Every sibling test in this
+        module mocks ``Extractor.register_*`` away, so none of them reaches the
+        validation at :file:`pcapkit/foundation/extraction.py` -- which is why the
+        defect survived: all three entry points rejected pcapkit's *own* engines,
+        reassembly and flow-tracing classes with
+        :exc:`~pcapkit.utilities.exceptions.RegistryError`.
+
+        The cause was that each check named the auto-registering public class
+        (``Engine``, ``Reassembly``, ``TraceFlow``) while every built-in subclasses
+        the ``*Base`` variant -- imported under an alias, e.g.
+        ``from ...engine import EngineBase as Engine`` at
+        :file:`pcapkit/foundation/engines/pcap.py`, specifically so that it is not
+        auto-registered. So ``issubclass(PCAP, Engine)`` was :data:`False` while
+        ``issubclass(PCAP, EngineBase)`` was :data:`True`.
+
+        Deliberately does **not** mock, because the point is to exercise the gate.
+
+        """
+        # NOTE: imported inside the test because ``setUp`` purges ``pcapkit`` from
+        # ``sys.modules``, which is why every sibling test imports locally too.
+        # ``Extractor`` is needed by name here so the registries can be read back.
+        from pcapkit.foundation.engines.pcap import PCAP as PCAP_Engine
+        from pcapkit.foundation.extraction import Extractor
+        from pcapkit.foundation.reassembly.ipv4 import IPv4 as IPv4_Reassembly
+        from pcapkit.foundation.registry.foundation import (register_extractor_engine,
+                                                            register_extractor_reassembly,
+                                                            register_extractor_traceflow)
+        from pcapkit.foundation.traceflow.tcp import TCP as TCP_TraceFlow
+        from pcapkit.utilities.exceptions import RegistryError
+
+        for name, func, klass, store in (
+            ('engine', register_extractor_engine, PCAP_Engine,
+             Extractor.__engine__),
+            ('reassembly', register_extractor_reassembly, IPv4_Reassembly,
+             Extractor.__reassembly__),
+            ('traceflow', register_extractor_traceflow, TCP_TraceFlow,
+             Extractor.__traceflow__),
+        ):
+            with self.subTest(kind=name, accepted=True):
+                # NOTE: a distinct key per call, so this neither collides with the
+                # built-in registrations already present nor emits the
+                # ``already registered, overwriting`` warning.
+                key = f'unit-513-{name}'
+                func(key, klass)
+
+                # NOTE: the registry is read back rather than merely asserting that
+                # no exception escaped. Without this, the test passes against a
+                # helper that runs the ``issubclass`` gate and then silently drops
+                # the class -- measured, not hypothetical: inserting a bare
+                # ``return`` after the check and before the registry write leaves
+                # this test reporting ``1 passed, 6 subtests passed``, exit 0.
+                self.assertIn(key, store)
+                self.assertIs(store[key], klass)
+
+        # And the gate still rejects something that is genuinely not a subclass --
+        # the widening must not have turned the check into a no-op.
+        for name, func in (
+            ('engine', register_extractor_engine),
+            ('reassembly', register_extractor_reassembly),
+            ('traceflow', register_extractor_traceflow),
+        ):
+            with self.subTest(kind=name, accepted=False):
+                with self.assertRaises(RegistryError):
+                    func(f'unit-513-reject-{name}', int)  # type: ignore[arg-type]
+
 
 if __name__ == '__main__':
     unittest.main()