Make the visitor class user customisable 🤦‍

Author: domdfcoding

flake8_helper/__init__.py

@@ -29,7 +29,7 @@
 # stdlib
 import ast
 from abc import ABC, abstractmethod
-from typing import Iterator, List, Tuple, Type, TypeVar
+from typing import Generic, Iterator, List, Tuple, Type, TypeVar
 
 __all__ = ["Visitor", "_P", "Plugin"]
 
@@ -40,6 +40,7 @@
 __email__: str = "dominic@davis-foster.co.uk"
 
 _P = TypeVar("_P", bound="Plugin")
+_V = TypeVar("_V", bound="Visitor")
 
 
 class Visitor(ast.NodeVisitor):
@@ -66,7 +67,7 @@ def report_error(self, node: ast.AST, error: str):
 				))
 
 
-class Plugin(ABC):
+class Plugin(ABC, Generic[_V]):
 	"""
 	Abstract base class for Flake8 plugins.
 
@@ -110,14 +111,23 @@ def version(self) -> str:
 
 		raise NotImplementedError
 
+	@property
+	@abstractmethod
+	def visitor_class(self) -> Type[_V]:
+		"""
+		The visitor class to use to traverse the AST.
+		"""
+
+		raise NotImplementedError
+
 	def run(self: _P) -> Iterator[Tuple[int, int, str, Type[_P]]]:
 		"""
 		Traverse the Abstract Syntax Tree and identify errors.
 
 		Yields a tuple of (line number, column offset, error message, type(self))
 		"""
 
-		visitor = Visitor()
+		visitor: _V = self.visitor_class()
 		visitor.visit(self._tree)
 
 		for line, col, msg in visitor.errors: