From 98f06e754e37cfe2fc3fb56ca782d8ef2c63c12f Mon Sep 17 00:00:00 2001
From: Frank Sachsenheim <funkyfuture@riseup.net>
Date: Mon, 23 Sep 2024 13:12:16 +0200
Subject: [PATCH] Improves annotations of PluginManager and contributed XPath
 functions

---
 _delb/plugins/__init__.py | 34 ++++++++++++++++++++++++----------
 _delb/typing.py           | 10 ++++++++--
 _delb/xpath/functions.py  | 10 +++++-----
 3 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/_delb/plugins/__init__.py b/_delb/plugins/__init__.py
index 2e597699..70b9c8d1 100644
--- a/_delb/plugins/__init__.py
+++ b/_delb/plugins/__init__.py
@@ -16,8 +16,8 @@
 from __future__ import annotations
 
 import sys
-from collections.abc import Callable, Iterable
-from typing import TYPE_CHECKING, Any
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, overload, Any
 
 if sys.version_info < (3, 10):  # DROPWITH Python3.9
     from importlib_metadata import entry_points
@@ -27,8 +27,14 @@
 
 if TYPE_CHECKING:
     from types import SimpleNamespace
-
-    from _delb.typing import Loader, LoaderConstraint
+    from delb import Document
+    from _delb.typing import (
+        GenericDecorated,
+        Loader,
+        LoaderConstraint,
+        SecondOrderDecorator,
+        XPathFunction,
+    )
 
 
 class DocumentMixinBase:
@@ -93,10 +99,10 @@ def _init_config(cls, config: SimpleNamespace, kwargs: dict[str, Any]):
 
 class PluginManager:
     def __init__(self):
-        self.document_mixins: list[type] = []
-        self.document_subclasses: list[type] = []
+        self.document_mixins: list[type[DocumentMixinBase]] = []
+        self.document_subclasses: list[type[Document]] = []
         self.loaders: list[Loader] = []
-        self.xpath_functions: dict[str, Callable] = {}
+        self.xpath_functions: dict[str, XPathFunction] = {}
 
     @staticmethod
     def load_plugins():
@@ -108,7 +114,7 @@ def load_plugins():
 
     def register_loader(
         self, before: LoaderConstraint = None, after: LoaderConstraint = None
-    ) -> Callable:
+    ) -> SecondOrderDecorator:
         """
         Registers a document loader.
 
@@ -200,7 +206,15 @@ def registrar(loader: Loader) -> Loader:
 
         return registrar
 
-    def register_xpath_function(self, arg: Callable | str) -> Callable:
+    @overload
+    def register_xpath_function(self, arg: str) -> SecondOrderDecorator: ...
+
+    @overload
+    def register_xpath_function(self, arg: GenericDecorated) -> GenericDecorated: ...
+
+    def register_xpath_function(
+        self, arg: str | GenericDecorated
+    ) -> SecondOrderDecorator | GenericDecorated:
         """
         Custom XPath functions can be defined as shown in the following example. The
         first argument to a function is always an instance of
@@ -231,7 +245,7 @@ def lowercase(_, string: str) -> str:
         """
         if isinstance(arg, str):
 
-            def wrapper(func):
+            def wrapper(func: XPathFunction) -> XPathFunction:
                 self.xpath_functions[arg] = func
                 return func
 
diff --git a/_delb/typing.py b/_delb/typing.py
index 2ea598ac..3d677014 100644
--- a/_delb/typing.py
+++ b/_delb/typing.py
@@ -18,11 +18,11 @@
 
 import sys
 
-from typing import TYPE_CHECKING, Any, Mapping, Union  # noqa: UNT001
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Union, TypeVar  # noqa: UNT001
 
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable
+    from collections.abc import Iterable
     from types import SimpleNamespace
 
     if sys.version_info < (3, 10):  # DROPWITH Python 3.9
@@ -33,6 +33,7 @@
     from lxml import etree
 
     from _delb.nodes import NodeBase, _TagDefinition
+    from _delb.xpath.ast import EvaluationContext
 
 
 if sys.version_info < (3, 11):  # DROPWITH Python 3.10
@@ -41,6 +42,9 @@
     from typing import Self
 
 
+GenericDecorated = TypeVar("GenericDecorated", bound=Callable[..., Any])
+SecondOrderDecorator: TypeAlias = "Callable[[GenericDecorated], GenericDecorated]"
+
 Filter: TypeAlias = "Callable[[NodeBase], bool]"
 NamespaceDeclarations: TypeAlias = Mapping[str | None, str]
 NodeSource: TypeAlias = Union[str, "NodeBase", "_TagDefinition"]
@@ -49,6 +53,8 @@
 Loader: TypeAlias = "Callable[[Any, SimpleNamespace], LoaderResult]"
 LoaderConstraint: TypeAlias = Union[Loader, "Iterable[Loader]", None]
 
+XPathFunction: TypeAlias = "Callable[[EvaluationContext, *Any], Any]"
+
 
 __all__ = (
     "Filter",
diff --git a/_delb/xpath/functions.py b/_delb/xpath/functions.py
index 1a50d276..d0e6f4f3 100644
--- a/_delb/xpath/functions.py
+++ b/_delb/xpath/functions.py
@@ -26,17 +26,17 @@
 
 
 @plugin_manager.register_xpath_function
-def concat(_, *strings: str) -> str:
+def concat(_: EvaluationContext, *strings: str) -> str:
     return "".join(strings)
 
 
 @plugin_manager.register_xpath_function
-def contains(_, string: str, substring: str) -> bool:
+def contains(_: EvaluationContext, string: str, substring: str) -> bool:
     return substring in string
 
 
 @plugin_manager.register_xpath_function
-def boolean(_, value: Any) -> bool:
+def boolean(_: EvaluationContext, value: Any) -> bool:
     return bool(value)
 
 
@@ -46,7 +46,7 @@ def last(context: EvaluationContext) -> int:
 
 
 @plugin_manager.register_xpath_function("not")
-def _not(_, value: Any) -> bool:
+def _not(_: EvaluationContext, value: Any) -> bool:
     return not value
 
 
@@ -56,7 +56,7 @@ def position(context: EvaluationContext) -> int:
 
 
 @plugin_manager.register_xpath_function("starts-with")
-def starts_with(_, string: str, prefix: str) -> bool:
+def starts_with(_: EvaluationContext, string: str, prefix: str) -> bool:
     return string.startswith(prefix)