如何通过Protocol类型注解限制Python泛型包装类的操作符支持范围？

Great question! This is a common pitfall when combining generic classes, operator overloading, and Python's type system. Let's break down why your current approach isn't working, then walk through the correct way to implement this.

Why Your Current __lt__ Implementation Fails #

Your initial attempt to annotate self: Wrap[Ord] doesn't work as expected because mypy doesn't link the method's Ord type variable to the class's generic T parameter.

When you write def __lt__(self: Wrap[Ord], other: Wrap[Ord]) -> bool, mypy interprets this as:

"This method accepts any Wrap instance where the wrapped type is some Ordered subtype, and any Wrap[Ord] as other."

It doesn't enforce that the specific Wrap instance calling the method has a T that implements Ordered. So even Wrap[Bar] or Wrap[lambda] can call __lt__ without type errors—since mypy doesn't check if their T matches Ord.

In contrast, annotating a class method parameter as Wrap[Ord] works because mypy checks the incoming argument's type directly, ensuring its T is Ordered.

The Fix: Conditional Method Overloading with Bounded Type Variables #

To make mypy correctly enforce that Wrap[T] only supports < when T implements Ordered, you need to explicitly link the method's constraint to the class's generic T. Here's how:

from typing import Generic, Protocol, TypeVar, Self

class Ordered(Protocol):
    """Types that support < comparison."""
    def __lt__(self, other: "Ordered") -> bool:
        """Determine if self is < other."""
        ...

# Define a private TypeVar bound to Ordered for method-level constraints
_OrdT = TypeVar('_OrdT', bound=Ordered)
T = TypeVar('T')

class Wrap(Generic[T]):
    """Wrap a value of type T."""
    val: T
    
    def __init__(self, val: T) -> None:
        """Wrap a value."""
        self.val = val

    def __lt__(self: Self[_OrdT], other: Wrap[_OrdT]) -> bool:
        """Compare wrapped values, only available when T is Ordered."""
        return self.val < other.val

What This Does: #

1. _OrdT Type Variable: This private type variable is bounded to Ordered, so it only accepts types that implement the < operator.
2. Self[_OrdT] Annotation: By annotating self with Self[_OrdT], we tell mypy that this __lt__ method only exists when the Wrap instance's T is a subtype of Ordered.
3. Consistent other Type: The other parameter uses Wrap[_OrdT] to ensure we're comparing against a wrapped value of the same ordered type (avoids comparing Wrap[int] to Wrap[str], for example).

Testing the Fix: #

• Valid usage (type-checks pass):

  w_int1 = Wrap(1)
  w_int2 = Wrap(2)
  print(w_int1 < w_int2)  # ✅ Mypy approves
  

• Invalid usage (triggers type error):

  class Bar:
      pass
  
  w_bar = Wrap(Bar())
  print(w_bar < w_bar)  # ❌ Mypy error: "Wrap[Bar]" has no attribute "__lt__"
  
  w_lambda = Wrap(lambda x: x)
  print(w_lambda < w_lambda)  # ❌ Mypy error: "Wrap[Callable[[Any], Any]]" has no attribute "__lt__"
  

Alternative: Using a Protocol for Wrapped Ordered Types #

Another approach is to define a protocol that describes wrapped ordered types, and let mypy automatically infer when Wrap[T] satisfies this protocol:

from typing import Generic, Protocol, TypeVar

class Ordered(Protocol):
    def __lt__(self, other: "Ordered") -> bool: ...

class WrappedOrdered(Protocol):
    val: Ordered
    def __lt__(self, other: "WrappedOrdered") -> bool: ...

T = TypeVar('T')

class Wrap(Generic[T]):
    val: T
    
    def __init__(self, val: T) -> None:
        self.val = val

    def __lt__(self, other: Wrap[T]) -> bool:
        return self.val < other.val

Mypy will automatically recognize that Wrap[T] implements WrappedOrdered whenever T implements Ordered. This works well if you need to pass wrapped ordered types to functions expecting WrappedOrdered, but the first approach is more direct for operator overloading.

Key Takeaway #

To restrict a generic class's methods to only work when its type parameter meets a protocol:

• Use a method-level type variable bounded to the protocol.
• Annotate self (or Self) with the generic class instantiated with this bounded type variable.
• This links the class's generic parameter to the method's constraint, making mypy enforce the correct type checks.

内容的提问来源于stack exchange，提问作者Thomas Mailund