Modifying rich comparison TypeError exception message for Generic Types-CodePudding

I have the following code where I create a generic type that supports comparison. When I compare different types, the code yields an exception (as expected). However, I want to modify the exception's message to make it more transparent.

Here is the code as it stands.

import abc
import typing

class Comparable(typing.Protocol):
    """A simple protocol to signal to TypeVar that each 
        value will have a less than (lt) dunder available."""

    @abc.abstractmethod
    def __lt__(self, other: typing.Any, /) -> bool:
        raise NotImplementedError

# each value that will be assigned to T will be 'Comparable', 
# i.e., meets Comparable's interface
T = typing.TypeVar("T", bound=Comparable)


class Node(typing.Generic[T]):
    """Acts as a wrapper around any value.
    This is to show in code the issue I am trying to fix"""

    def __init__(self, value: T) -> None:
        self.value = value

    def __lt__(self, __other: typing.Union[T, Node[T]]) -> bool:
        """Implements support for the '<' operator"""
        try:
            if isinstance(__other, Node):
                return self.value < __other.value
            return self.value < __other
        except TypeError:
            return NotImplemented

The code above works as expected, and MyPy is happy. Types are inferred when an instance of Node is created with some value, and Node[type] can be used to annotate, as expected.

Here is some example of using Node and the issue I am facing.

value = Node(1)  # value: Node[int] = Node(1) -- value has a type of Node[int]
value2 = Node(2)  # likewise value2 has a type of Node[int]

# Example 1
print(
    value < 1
)  # -> False; the less than (lt) dunder can deal with Node[int] < int. 
# As you recall, __other is meant to accept T and Node[T]. 
# In this case, __other is 1, an int which is T.

# Example 2
print(
    value < value2
)  # -> True; the less than (lt) dunder should be able to deal with 
# Node[int] < Node[int] as __other would be Node[T]


# consider this
print(
    value < "0"
)  # As expected, this will fail because we cannot compare int and str; 
# likewise, we can't compare Node[int] with Node[str].
# Yields; <<Exeption>> 
# TypeError: '<' not supported between instances of 'Node' and 'str'

I am not sure if this is possible; however, I want to modify the output for the following exception such that it prints:

TypeError: '<' not supported between instances of 'Node[int]' and 'str'

Because technically < is supported between Node and str.

CodePudding user response：

I'm afraid you are out of luck, if you want to have that exact same error message but with the added type argument. The logic that handles the error message that you get, when a rich comparison method returns NotImplemented is an implementation detail of Python itself. You can for example see, how the error message is formed in the CPython 3.11.1 implementation of do_richcompare.

The closest you can get is modifying the error message and re-raising the TypeError manually, but it will include the entire stack trace including the line(s) from your __lt__ method.

Here is a full working example:

from __future__ import annotations
from typing import Any, Generic, Protocol, TypeVar, Union, cast


class Comparable(Protocol):
    def __lt__(self, other: Any, /) -> bool:
        ...


T = TypeVar("T", bound=Comparable)


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

    def __lt__(self, __other: Union[T, Node[T]]) -> bool:
        try:
            if isinstance(__other, Node):
                return self.value < __other.value
            return self.value < __other
        except TypeError as exc:
            # Modify exception message to to specify Node[T]
            # and raise the TypeError manually
            cls_name = self.value.__class__.__name__
            msg = cast(str, exc.args[0])
            msg = msg.replace(f"{cls_name}", f"Node[{cls_name}]", 1)
            if isinstance(__other, Node):
                cls_name = __other.value.__class__.__name__
                msg = msg.replace(
                    f"and '{cls_name}'",
                    f"and 'Node[{cls_name}]'",
                )
            exc.args = (msg, )   exc.args[1:]
            raise exc


def test() -> None:
    print(Node(1) < 1)
    print(Node(1) < Node(2))
    # Node(1) < "0"
    # Node(1) < Node("0")


if __name__ == "__main__":
    test()

The output like this is still False \ True.

If you uncomment Node(1) < "0", you'll get something like this:

Traceback (most recent call last):
  File "[...].py", line 47, in <module>
    test()
  File "[...].py", line 42, in test
    Node(1) < "0"
  File "[...].py", line 36, in __lt__
    raise exc from None
  File "[...].py", line 22, in __lt__
    return self.value < __other
TypeError: '<' not supported between instances of 'Node[int]' and 'str'

If you instead uncomment the Node(1) < Node("0"), you'll get this:

Traceback (most recent call last):
  File "[...].py", line 47, in <module>
    test()
  File "[...].py", line 43, in test
    Node(1) < Node("0")
  File "[...].py", line 36, in __lt__
    raise exc from None
  File "[...].py", line 21, in __lt__
    return self.value < __other.value
TypeError: '<' not supported between instances of 'Node[int]' and 'Node[str]'

I suppose, if instead of modifying the message on the existing exeption instance, you raised an entirely new instance of TypeError and added from None, you could cut off one additional step in the stack trace message, but that could also cost you some potentially useful additional exception details in some situations.

On an unrelated note, there is no need for abc.abstractmethod on Comparable.__lt__ because a Protocol may never be instantiated anyway. The method also doesn't need a body.