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) |
Update on 2019/03/06:
- Add strict type checking
- Force cast from mapping
example:@data class SecondLevel: foo: int bar: str @data class TopLevel: second: SecondLevel baz: float # This TopLevel({"foo": 12, "bar": "some string"}, 0.6) # equivalent to TopLevel(SecondLevel(12, "some string"), 0.6)
- Add default values
example:@data class Foo: bar: int = 5 baz: str = "some string" hoge: int # This Foo(12) # equivalent to Foo(12, 5) # or Foo(12, 5, "some string")
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Usage:
stdout:
This decorator automatically implements methods (__init__, __eq__, __ne__, __iter__, __str__, __repr__) to a class with annotation which is decorated. Target class must have least one member annotation.
If you need a variable of which type is Any, annotate with
None
.e.g.
If you use this decorator, you can enjoy more and more happy, easy, simple and clear Python programming!