datatype.py

"""
DataType
--------

A NamedTuple like class except it also defines strict type dependent
__eq__, __ne__  and __hash__. This ensures that instances can be compared,
hashed in dict, sets, ... I.e.

Example
-------

>>> class A(DataType):
>>>     a: str
>>>     b: int
>>>
>>>
>>> class B(DataType):
>>>     a: str
>>>     b: int
>>>
>>>
>>> assert A("a", 1) != B("a", 1)
>>> assert A("a", 1) != ("a", 1)
>>> assert not A("a", 1) == B("a", 1)
>>> assert hash(A("a", 1)) != hash(B("a", 1))
>>> assert len({A("a", 1), B("a", 1), ("a", 1)}) == 3

"""
from typing import NamedTuple
from typing import NamedTuple as DataType

__all__ = ["DataType"]


_NamedTupleMeta = type(NamedTuple)  # type: ignore


class _DataTypeMethods:
    def __eq__(self: tuple, other):  # type: ignore
        """Equal if `other` has the same type and all elements compare equal."""
        if type(self) is not type(other):
            return False
        return tuple.__eq__(self, other)  # type: ignore

    def __ne__(self: tuple, other):  # type: ignore
        return not self == other

    def __hash__(self: tuple):   # type: ignore
        return hash((type(self), tuple.__hash__(self)))  # type: ignore


class _DataTypeMeta(_NamedTupleMeta):  # type: ignore
    def __new__(cls, typename, bases, ns, **kwargs):
        if ns.get('_root', False):
            return super().__new__(cls, typename, bases, ns)
        cls = super().__new__(cls, typename, bases, ns, **kwargs)
        cls.__eq__ = _DataTypeMethods.__eq__
        cls.__ne__ = _DataTypeMethods.__ne__
        cls.__hash__ = _DataTypeMethods.__hash__
        cls.__module__ = ns.get("__module__", cls.__module__)
        return cls


# Replace the DataType (alias for NamedTuple) in globals without the type
# checker being any the wiser. NamedTuple are special cased. The only way
# for type checkers to consistently apply NamedTuple aliasing is
# import ... as ...,
globals()["DataType"] = _DataTypeMeta("DataType", (), {"_root": True})



class A(DataType):
    a: str
    b: int


class B(DataType):
    a: str
    b: int


assert A("a", 1) != B("a", 1)
assert A("a", 1) != ("a", 1)
assert not A("a", 1) == B("a", 1)
assert hash(A("a", 1)) != hash(B("a", 1))
assert len({A("a", 1), B("a", 1), ("a", 1)}) == 3