Last active
March 6, 2019 13:02
-
-
Save frodo821/bf4d69d42f06b9735813176d1ab87873 to your computer and use it in GitHub Desktop.
Data class decorator for Python!
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
from sys import _getframe as gf | |
class DataClass: | |
pass | |
def data(cls): | |
""" | |
Data class annotation | |
Usage: | |
```Python | |
>>> @data class Foo: | |
... bar: int | |
... baz: str | |
... | |
>>> dat = Foo(12, "test") | |
>>> dat | |
Foo(a=12, b="test") | |
>>> a, b = dat | |
>>> a, b | |
(12, "test") | |
>>> dat == Foo(12, "test") | |
True | |
>>> dat == Foo(11, "test") | |
False | |
>>> dat != Foo(12, "none") | |
True | |
``` | |
""" | |
if not hasattr(cls, '__annotations__'): | |
raise TypeError("Data class must have least one type annotation.") | |
lcls = dict(getattr(cls, "__dict__", {})) | |
ann = cls.__annotations__ | |
lcls['__annotations__'] = ann | |
lcls['__slots__'] = list(ann.keys()) | |
init = [] | |
ibody = [] | |
ebody = [] | |
nbody = [] | |
tbody = [] | |
sbody = [] | |
for k, v in ann.items(): | |
default = repr(lcls.pop(k)) if k in lcls else '' | |
if isinstance(v, type): | |
init.append((f"{k}: {v.__name__}{'='+default if default else ''}", bool(default))) | |
ibody.append(f"if isinstance({k}, {v.__name__}):") | |
ibody.append(f" self.{k} = {k}") | |
ibody.append(f"else:") | |
if issubclass(v, DataClass): | |
ibody.append(f" self.{k} = {v.__name__}(**{k})") | |
else: | |
ibody.append( | |
f" raise TypeError('Unexpected instance of type \'{{type({k}).__name__}}\'')") | |
else: | |
init.append((f"{k}{'='+default if default else ''}", bool(default))) | |
ibody.append(f"self.{k} = {k}") | |
ebody.append(f"self.{k} == other.{k}") | |
nbody.append(f"self.{k} != other.{k}") | |
tbody.append(f"yield self.{k}") | |
sbody.append(f"{k}={{repr(self.{k})}}") | |
init.sort(key=lambda x: x[1]) | |
init = map(lambda x: x[0], init) | |
src = (f'def __init__(self, {",".join(init)}):\n' | |
f'{f"{chr(10)} ".join(ibody)}\n' | |
f'def __eq__(self, other): return {" and ".join(ebody)}\n' | |
f'def __ne__(self, other): return {" or ".join(nbody)}\n' | |
f'def __str__(self): return f"{cls.__name__}({", ".join(sbody)})"\n' | |
'def __repr__(self): return self.__str__()\n' | |
'def __iter__(self):\n' | |
f' {f"{chr(10)} ".join(tbody)}') | |
gns = gf(1).f_globals | |
gns['DataClass'] = DataClass | |
exec(src, gns, lcls) | |
return type(cls.__name__, (DataClass,) + tuple(type.mro(cls)[1:]), lcls) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Update on 2019/03/06:
example:
example: