Extension framework optimization

src/markitect_tool/extension/registry.py

@@ -112,6 +112,8 @@ class ExtensionRegistry:
 
     def __init__(self, descriptors: Iterable[ExtensionDescriptor] | None = None) -> None:
         self._descriptors: dict[str, ExtensionDescriptor] = {}
+        self._by_kind: dict[str, set[str]] = {}
+        self._by_capability: dict[str, set[str]] = {}
         for descriptor in descriptors or []:
             self.register(descriptor)
 
@@ -119,6 +121,9 @@ class ExtensionRegistry:
         if descriptor.id in self._descriptors:
             raise ExtensionRegistryError(f"Duplicate extension id `{descriptor.id}`")
         self._descriptors[descriptor.id] = descriptor
+        self._by_kind.setdefault(descriptor.kind, set()).add(descriptor.id)
+        for capability in descriptor.capabilities:
+            self._by_capability.setdefault(capability.id, set()).add(descriptor.id)
 
     def get(self, extension_id: str) -> ExtensionDescriptor:
         try:
@@ -127,16 +132,16 @@ class ExtensionRegistry:
             raise ExtensionRegistryError(f"Unknown extension `{extension_id}`") from exc
 
     def list(self, *, kind: str | None = None) -> list[ExtensionDescriptor]:
-        descriptors = [self._descriptors[key] for key in sorted(self._descriptors)]
         if kind is None:
-            return descriptors
-        return [descriptor for descriptor in descriptors if descriptor.kind == kind]
+            ids = sorted(self._descriptors)
+        else:
+            ids = sorted(self._by_kind.get(kind, set()))
+        return [self._descriptors[key] for key in ids]
 
     def require_capability(self, capability_id: str) -> list[ExtensionDescriptor]:
         return [
-            descriptor
-            for descriptor in self.list()
-            if any(capability.id == capability_id for capability in descriptor.capabilities)
+            self._descriptors[extension_id]
+            for extension_id in sorted(self._by_capability.get(capability_id, set()))
         ]
 
     def check_dependencies(