Last active
July 28, 2020 14:22
-
-
Save ales-erjavec/822c8d27cbd1b1c263e9af043e96c02f to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment