nteract / scrapbook

A library for recording and reading data in notebooks.
https://nteract-scrapbook.readthedocs.io
BSD 3-Clause "New" or "Revised" License
281 stars 26 forks source link

scrap.Scrap as a dataclass #89

Open philippefutureboy opened 3 years ago

philippefutureboy commented 3 years ago

Hey there!

While I was fiddling around with scrapbook, I ended up writing a dataclass for scrap.Scrap items. I noticed in the source that you have a comment that says:

# dataclasses would be nice here...
Scrap = namedtuple("Scrap", ["name", "data", "encoder", "display"])
Scrap.__new__.__defaults__ = (None,)

So I figured you could be interested by the implementation I made! :) See below:

@dataclass
class Scrap(collections.abc.Mapping):
    name: str
    data: Any
    encoder: str
    display: str = None

    def keys(self):
        return self.__dataclass_fields__.keys() # pylint: disable=no-member

    def __getitem__(self, key: str) -> Any:
        if isinstance(key, str):
            if not hasattr(self, key):
                raise KeyError(key)
            return getattr(self, key)
        else:
            raise TypeError(f"Unsupported key type: {type(key)}")

    def __len__(self):
        return len(self.__dataclass_fields__.keys()) # pylint: disable=no-member

    def __iter__(self):
        return (attr for attr in dataclasses.astuple(self))

    def asdict(self) -> dict:
        return dataclasses.asdict(self)

    def astuple(self) -> tuple:
        return dataclasses.astuple(self)