From f2cb15d0e216003f8d3cee9f392b1c9f5ba566f9 Mon Sep 17 00:00:00 2001
From: thatmattlove <matt@hyperglass.dev>
Date: Mon, 13 Sep 2021 14:10:32 -0700
Subject: [PATCH] Fix typing issues

---
 hyperglass/execution/drivers/_common.py |  2 +-
 hyperglass/execution/main.py            |  2 +-
 hyperglass/plugins/_base.py             | 16 ++++++++++++++--
 hyperglass/plugins/_output.py           |  4 ++--
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/hyperglass/execution/drivers/_common.py b/hyperglass/execution/drivers/_common.py
index bc4c4ae..085af2d 100644
--- a/hyperglass/execution/drivers/_common.py
+++ b/hyperglass/execution/drivers/_common.py
@@ -37,7 +37,7 @@ class Connection(ABC):
         """Return a preconfigured sshtunnel.SSHTunnelForwarder instance."""
         pass
 
-    async def response(self, output: Sequence[str]) -> Union[OutputDataModel, str]:
+    async def response(self, output: Sequence[str]) -> Union["OutputDataModel", str]:
         """Send output through common parsers."""
 
         log.debug("Pre-parsed responses:\n{}", output)
diff --git a/hyperglass/execution/main.py b/hyperglass/execution/main.py
index 91a2bbf..8d49d74 100644
--- a/hyperglass/execution/main.py
+++ b/hyperglass/execution/main.py
@@ -45,7 +45,7 @@ def handle_timeout(**exc_args: Any) -> Callable:
     return handler
 
 
-async def execute(query: "Query") -> Union[OutputDataModel, str]:
+async def execute(query: "Query") -> Union["OutputDataModel", str]:
     """Initiate query validation and execution."""
 
     output = params.messages.general
diff --git a/hyperglass/plugins/_base.py b/hyperglass/plugins/_base.py
index 977bf2f..a3b782b 100644
--- a/hyperglass/plugins/_base.py
+++ b/hyperglass/plugins/_base.py
@@ -54,7 +54,19 @@ class HyperglassPlugin(BaseModel, ABC):
         super().__init__(name=name, **kwargs)
 
 
-class DirectivePlugin(HyperglassPlugin):
-    """Plugin associated with directives."""
+class DirectivePlugin(BaseModel):
+    """Plugin associated with directives.
+
+    Should always be subclassed with `HyperglassPlugin`.
+    """
 
     directives: Sequence[str] = ()
+
+
+class DeviceTypePlugin(BaseModel):
+    """Plugin associated with specific device types.
+
+    Should always be subclassed with `HyperglassPlugin`.
+    """
+
+    device_types: Sequence[str] = ()
diff --git a/hyperglass/plugins/_output.py b/hyperglass/plugins/_output.py
index 8962b14..a30d2e1 100644
--- a/hyperglass/plugins/_output.py
+++ b/hyperglass/plugins/_output.py
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Union, Sequence
 from hyperglass.log import log
 
 # Local
-from ._base import DirectivePlugin
+from ._base import DirectivePlugin, DeviceTypePlugin, HyperglassPlugin
 
 if TYPE_CHECKING:
     # Project
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 OutputType = Union["OutputDataModel", Sequence[str]]
 
 
-class OutputPlugin(DirectivePlugin):
+class OutputPlugin(HyperglassPlugin, DirectivePlugin, DeviceTypePlugin):
     """Plugin to interact with device command output."""
 
     def process(self, output: OutputType, device: "Device") -> OutputType: