Add dictionary mapping command_id to its associated MenuCommands for each plugin

src/npe2/_plugin_manager.py

@@ -4,7 +4,7 @@
 import os
 import urllib
 import warnings
-from collections import Counter
+from collections import Counter, defaultdict
 from fnmatch import fnmatch
 from importlib import metadata
 from logging import getLogger
@@ -38,6 +38,7 @@
 if TYPE_CHECKING:
     from .manifest.contributions import (
         CommandContribution,
+        MenuCommand,
         MenuItem,
         ReaderContribution,
         SampleDataContribution,
@@ -235,6 +236,9 @@ def __init__(
         self._manifests: Dict[PluginName, PluginManifest] = {}
         self.events = PluginManagerEvents(self)
         self._npe1_adapters: List[NPE1Adapter] = []
+        self._command_menu_map: Dict[
+            str, Dict[str, Dict[str, List[MenuCommand]]]
+        ] = defaultdict(dict)
 
         # up to napari 0.4.15, discovery happened in the init here
         # so if we're running on an older version of napari, we need to discover
@@ -358,14 +362,36 @@ def register(
             self._npe1_adapters.append(manifest)
         else:
             self._contrib.index_contributions(manifest)
+            self._populate_command_menu_map(manifest)
         self.events.plugins_registered.emit({manifest})
 
+    def _populate_command_menu_map(self, manifest: PluginManifest):
+        self._command_menu_map[manifest.name] = defaultdict(dict)
+        for command in manifest.contributions.commands or ():
+            # command IDs are keys in map
+            # each value is a dict menu_id: list of MenuCommands
+            # for the command and menu
+            associated = self._get_associated_menus(manifest, command.id)
+            self._command_menu_map[manifest.name][command.id] = associated
+
+    def _get_associated_menus(
+        self, manifest: PluginManifest, command_id: str
+    ) -> Dict[str, List[MenuCommand]]:
+        menus = manifest.contributions.menus or dict()
+        associated_menus = defaultdict(list)
+        for menu_id, items in menus.items():
+            for item in items:
+                if getattr(item, "command", "") == command_id:
+                    associated_menus[menu_id].append(item)
+        return associated_menus
+
     def unregister(self, key: PluginName):
         """Unregister plugin named `key`."""
         if key not in self._manifests:
             raise ValueError(f"No registered plugin named {key!r}")  # pragma: no cover
         self.deactivate(key)
         self._contrib.remove_contributions(key)
+        self._command_menu_map.pop(key)
         self._manifests.pop(key)
 
     def activate(self, key: PluginName) -> PluginContext:
@@ -448,6 +474,7 @@ def enable(self, plugin_name: PluginName) -> None:
         mf = self._manifests.get(plugin_name)
         if mf is not None:
             self._contrib.index_contributions(mf)
+            self._populate_command_menu_map(mf)
         self.events.enablement_changed({plugin_name}, {})
 
     def disable(self, plugin_name: PluginName) -> None:
@@ -467,6 +494,7 @@ def disable(self, plugin_name: PluginName) -> None:
 
         self._disabled_plugins.add(plugin_name)
         self._contrib.remove_contributions(plugin_name)
+        self._command_menu_map.pop(plugin_name)
         self.events.enablement_changed({}, {plugin_name})
 
     def is_disabled(self, plugin_name: str) -> bool:

tests/test_plugin_manager.py

@@ -236,3 +236,33 @@ def dummy_error():
     pm.get_context("test").register_disposable(dummy_error)
     pm.deactivate("test")
     assert caplog.records[0].msg == "Error while disposing test; This is an error"
+
+
+def test_command_menu_map(uses_sample_plugin, plugin_manager: PluginManager):
+    """Test that the command menu map is correctly populated."""
+    pm = PluginManager.instance()
+    assert SAMPLE_PLUGIN_NAME in pm._manifests
+    assert SAMPLE_PLUGIN_NAME in pm._command_menu_map
+
+    # contains correct commands
+    command_menu_map = pm._command_menu_map[SAMPLE_PLUGIN_NAME]
+    assert "my-plugin.hello_world" in command_menu_map
+    assert "my-plugin.another_command" in command_menu_map
+
+    # commands point to correct menus
+    assert len(cmd_menu := command_menu_map["my-plugin.hello_world"]) == 1
+    assert "/napari/layer_context" in cmd_menu
+    assert len(cmd_menu := command_menu_map["my-plugin.another_command"]) == 1
+    assert "mysubmenu" in cmd_menu
+
+    # enable/disable
+    pm.disable(SAMPLE_PLUGIN_NAME)
+    assert SAMPLE_PLUGIN_NAME not in pm._command_menu_map
+    pm.enable(SAMPLE_PLUGIN_NAME)
+    assert SAMPLE_PLUGIN_NAME in pm._command_menu_map
+
+    # register/unregister
+    pm.unregister(SAMPLE_PLUGIN_NAME)
+    assert SAMPLE_PLUGIN_NAME not in pm._command_menu_map
+    pm.register(SAMPLE_PLUGIN_NAME)
+    assert SAMPLE_PLUGIN_NAME in pm._command_menu_map