use _SupportsComparison type from typeshed

Author: miguelgrinberg

elasticsearch_dsl/wrappers.py

@@ -17,56 +17,37 @@
 
 import operator
 from typing import (
-    Any,
+    TYPE_CHECKING,
     Callable,
     ClassVar,
     Dict,
     Literal,
     Mapping,
     Optional,
-    Protocol,
     Tuple,
     TypeVar,
     Union,
     cast,
 )
 
+if TYPE_CHECKING:
+    from _operator import _SupportsComparison
+
 from typing_extensions import TypeAlias
 
 from .utils import AttrDict
 
-
-class SupportsDunderLT(Protocol):
-    def __lt__(self, other: Any, /) -> Any: ...
-
-
-class SupportsDunderGT(Protocol):
-    def __gt__(self, other: Any, /) -> Any: ...
-
-
-class SupportsDunderLE(Protocol):
-    def __le__(self, other: Any, /) -> Any: ...
-
-
-class SupportsDunderGE(Protocol):
-    def __ge__(self, other: Any, /) -> Any: ...
-
-
-SupportsComparison: TypeAlias = Union[
-    SupportsDunderLE, SupportsDunderGE, SupportsDunderGT, SupportsDunderLT
-]
-
 ComparisonOperators: TypeAlias = Literal["lt", "lte", "gt", "gte"]
-RangeValT = TypeVar("RangeValT", bound=SupportsComparison)
+RangeValT = TypeVar("RangeValT", bound=_SupportsComparison)
 
-__all__ = ["Range", "SupportsComparison"]
+__all__ = ["Range"]
 
 
 class Range(AttrDict[ComparisonOperators, RangeValT]):
     OPS: ClassVar[
         Mapping[
             ComparisonOperators,
-            Callable[[SupportsComparison, SupportsComparison], bool],
+            Callable[[_SupportsComparison, _SupportsComparison], bool],
         ]
     ] = {
         "lt": operator.lt,
@@ -116,7 +97,7 @@ def __contains__(self, item: object) -> bool:
 
         # Cast to tell mypy whe have checked it and its ok to use the comparison methods
         # on `item`
-        item = cast(SupportsComparison, item)
+        item = cast(_SupportsComparison, item)
 
         for op in self.OPS:
             if op in self._d_ and not self.OPS[op](item, self._d_[op]):

tests/test_wrappers.py

@@ -16,12 +16,14 @@
 #  under the License.
 
 from datetime import datetime, timedelta
-from typing import Any, Mapping, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence
+
+if TYPE_CHECKING:
+    from _operator import _SupportsComparison
 
 import pytest
 
 from elasticsearch_dsl import Range
-from elasticsearch_dsl.wrappers import SupportsComparison
 
 
 @pytest.mark.parametrize(
@@ -37,7 +39,7 @@
     ],
 )
 def test_range_contains(
-    kwargs: Mapping[str, SupportsComparison], item: SupportsComparison
+    kwargs: Mapping[str, _SupportsComparison], item: _SupportsComparison
 ) -> None:
     assert item in Range(**kwargs)
 
@@ -53,7 +55,7 @@ def test_range_contains(
     ],
 )
 def test_range_not_contains(
-    kwargs: Mapping[str, SupportsComparison], item: SupportsComparison
+    kwargs: Mapping[str, _SupportsComparison], item: _SupportsComparison
 ) -> None:
     assert item not in Range(**kwargs)
 
@@ -69,7 +71,7 @@ def test_range_not_contains(
     ],
 )
 def test_range_raises_value_error_on_wrong_params(
-    args: Sequence[Any], kwargs: Mapping[str, SupportsComparison]
+    args: Sequence[Any], kwargs: Mapping[str, _SupportsComparison]
 ) -> None:
     with pytest.raises(ValueError):
         Range(*args, **kwargs)
@@ -85,8 +87,8 @@ def test_range_raises_value_error_on_wrong_params(
     ],
 )
 def test_range_lower(
-    range: Range[SupportsComparison],
-    lower: Optional[SupportsComparison],
+    range: Range[_SupportsComparison],
+    lower: Optional[_SupportsComparison],
     inclusive: bool,
 ) -> None:
     assert (lower, inclusive) == range.lower
@@ -102,8 +104,8 @@ def test_range_lower(
     ],
 )
 def test_range_upper(
-    range: Range[SupportsComparison],
-    upper: Optional[SupportsComparison],
+    range: Range[_SupportsComparison],
+    upper: Optional[_SupportsComparison],
     inclusive: bool,
 ) -> None:
     assert (upper, inclusive) == range.upper