Pass full query to output plugin instead of device

hyperglass/execution/drivers/_common.py

@@ -43,9 +43,7 @@ class Connection(ABC):
 
         log.debug("Pre-parsed responses:\n{}", output)
 
-        response = self.plugin_manager.execute(
-            directive=self.query_data.directive, output=output, device=self.device
-        )
+        response = self.plugin_manager.execute(output=output, query=self.query_data)
 
         if response is None:
             response = ()

hyperglass/plugins/_builtin/bgp_route_juniper.py

@@ -22,7 +22,7 @@ if TYPE_CHECKING:
 
     # Project
     from hyperglass.models.data import OutputDataModel
-    from hyperglass.models.config.devices import Device
+    from hyperglass.models.api.query import Query
 
     # Local
     from .._output import OutputType
@@ -127,14 +127,14 @@ class BGPRoutePluginJuniper(OutputPlugin):
         "__hyperglass_juniper_bgp_community_table__",
     )
 
-    def process(self, output: "OutputType", device: "Device") -> "OutputType":
+    def process(self, *, output: "OutputType", query: "Query") -> "OutputType":
         """Parse Juniper response if data is a string (and is therefore unparsed)."""
         should_process = all(
             (
                 isinstance(output, (list, tuple)),
-                device.platform in self.platforms,
-                device.structured_output is True,
-                device.has_directives(*self.directives),
+                query.device.platform in self.platforms,
+                query.device.structured_output is True,
+                query.device.has_directives(*self.directives),
             )
         )
         if should_process:

hyperglass/plugins/_builtin/remove_command.py

@@ -14,7 +14,7 @@ from .._output import OutputType, OutputPlugin
 
 if TYPE_CHECKING:
     # Project
-    from hyperglass.models.config.devices import Device
+    from hyperglass.models.api.query import Query
 
 
 class RemoveCommand(OutputPlugin):
@@ -22,13 +22,13 @@ class RemoveCommand(OutputPlugin):
 
     __hyperglass_builtin__: bool = PrivateAttr(True)
 
-    def process(self, device_output: OutputType, device: "Device") -> Sequence[str]:
+    def process(self, *, output: OutputType, query: "Query") -> Sequence[str]:
         """Remove anything before the command if found in output."""
 
         def _remove_command(output_in: str) -> str:
-            output_out = device_output.strip().split("\n")
+            output_out = output_in.strip().split("\n")
 
-            for command in device.directive_commands:
+            for command in query.device.directive_commands:
                 for line in output_out:
                     if command in line:
                         idx = output_out.index(line) + 1
@@ -36,7 +36,7 @@ class RemoveCommand(OutputPlugin):
 
             return "\n".join(output_out)
 
-        if is_series(device_output):
-            return tuple(_remove_command(o) for o in device_output)
+        if is_series(output):
+            return tuple(_remove_command(o) for o in output)
 
-        return device_output
+        return output

hyperglass/plugins/_manager.py

@@ -19,7 +19,6 @@ if t.TYPE_CHECKING:
     from hyperglass.state import HyperglassState
     from hyperglass.models.api.query import Query
     from hyperglass.models.directive import Directive
-    from hyperglass.models.config.devices import Device
 
 PluginT = t.TypeVar("PluginT", bound=HyperglassPlugin)
 
@@ -148,17 +147,17 @@ class InputPluginManager(PluginManager[InputPlugin], type="input"):
 class OutputPluginManager(PluginManager[OutputPlugin], type="output"):
     """Manage Output Processing Plugins."""
 
-    def execute(
-        self: "OutputPluginManager", *, directive: "Directive", output: OutputType, device: "Device"
-    ) -> OutputType:
+    def execute(self: "OutputPluginManager", *, output: OutputType, query: "Query") -> OutputType:
         """Execute all output parsing plugins.
 
         The result of each plugin is passed to the next plugin.
         """
         result = output
-        for plugin in (plugin for plugin in self.plugins if directive.id in plugin.directives):
+        for plugin in (
+            plugin for plugin in self.plugins if query.directive.id in plugin.directives
+        ):
             if result is False:
                 return result
             # Pass the result of each plugin to the next plugin.
-            result = plugin.process(result, device)
+            result = plugin.process(output=result, query=query)
         return result

hyperglass/plugins/_output.py

@@ -13,7 +13,7 @@ from ._base import PlatformPlugin, DirectivePlugin, HyperglassPlugin
 if TYPE_CHECKING:
     # Project
     from hyperglass.models.data import OutputDataModel
-    from hyperglass.models.config.devices import Device
+    from hyperglass.models.api.query import Query
 
 OutputType = Union["OutputDataModel", Series[str]]
 
@@ -21,7 +21,7 @@ OutputType = Union["OutputDataModel", Series[str]]
 class OutputPlugin(HyperglassPlugin, DirectivePlugin, PlatformPlugin):
     """Plugin to interact with device command output."""
 
-    def process(self, output: OutputType, device: "Device") -> OutputType:
+    def process(self, *, output: OutputType, query: "Query") -> OutputType:
         """Process or manipulate output from a device."""
         log.warning("Output plugin '{}' has not implemented a 'process()' method", self.name)
         return output