tool protocol

Author: KonstantAnxiety

lib/dl_formula_testing/dl_formula_testing/tool_runner.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import argparse
 from contextlib import (
     contextmanager,
     redirect_stderr,
@@ -11,6 +12,7 @@
     TYPE_CHECKING,
     Generator,
     Optional,
+    Protocol,
     TextIO,
 )
 
@@ -19,6 +21,12 @@
     from argparse import ArgumentParser
 
 
+class Tool(Protocol):
+    @classmethod
+    def run(cls, args: argparse.Namespace) -> None:
+        raise NotImplementedError()
+
+
 @contextmanager
 def redirect_stdin(stream: Optional[TextIO] = None) -> Generator[None, None, None]:
     if stream is None:
@@ -34,7 +42,7 @@ def redirect_stdin(stream: Optional[TextIO] = None) -> Generator[None, None, Non
 
 
 class ToolRunner:
-    def __init__(self, parser: ArgumentParser, tool_cls: type):
+    def __init__(self, parser: ArgumentParser, tool_cls: type[Tool]):
         self.parser = parser
         self.tool_cls = tool_cls