Refactor PrettyPrinter type-checking and extend it to all commands

Author: sara hartse

sdb/command.py

@@ -240,6 +240,29 @@ def _call(self,
         """
         raise NotImplementedError()
 
+    def __input_type_check(
+            self, objs: Iterable[drgn.Object]) -> Iterable[drgn.Object]:
+        assert self.input_type is not None
+        valid_types = [type_canonicalize_name(self.input_type)]
+
+        #
+        # Some commands support multiple input types. Check for InputHandler
+        # implementations to expand the set of valid input types.
+        #
+        for (_, method) in inspect.getmembers(self, inspect.ismethod):
+            if not hasattr(method, "input_typename_handled"):
+                continue
+            valid_types.append(
+                type_canonicalize_name(method.input_typename_handled))
+
+        for obj in objs:
+            if type_canonical_name(obj.type_) not in valid_types:
+                raise CommandError(
+                    self.name,
+                    f'exepected input of type {self.input_type}, but received '
+                    f'type {obj.type_}')
+            yield obj
+
     def __invalid_memory_objects_check(self, objs: Iterable[drgn.Object],
                                        fatal: bool) -> Iterable[drgn.Object]:
         """
@@ -277,7 +300,11 @@ def call(self, objs: Iterable[drgn.Object]) -> Iterable[drgn.Object]:
         # the command is running.
         #
         try:
-            result = self._call(objs)
+            if self.input_type and objs:
+                result = self._call(self.__input_type_check(objs))
+            else:
+                result = self._call(objs)
+
             if result is not None:
                 #
                 # The whole point of the SingleInputCommands are that
@@ -637,17 +664,7 @@ def _call(  # type: ignore[return]
         This function will call pretty_print() on each input object,
         verifying the types as we go.
         """
-
-        assert self.input_type is not None
-        type_name = type_canonicalize_name(self.input_type)
-        for obj in objs:
-            if type_canonical_name(obj.type_) != type_name:
-                raise CommandError(
-                    self.name,
-                    f'exepected input of type {self.input_type}, but received '
-                    f'type {obj.type_}')
-
-            self.pretty_print([obj])
+        self.pretty_print(objs)
 
 
 class Locator(Command):

tests/integration/data/regression_output/linux/addr spa_namespace_avl | cpu_counter_sum

@@ -1 +1 @@
-sdb: cpu_counter_sum: input is not a percpu_counter
+sdb: cpu_counter_sum: exepected input of type struct percpu_counter *, but received type avl_tree_t *