Skip to content

Commit

Permalink
[skip ci] Add PickleItem
Browse files Browse the repository at this point in the history
  • Loading branch information
thomass-dev committed Jan 7, 2025
1 parent c08f284 commit 06ff76f
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 7 deletions.
10 changes: 3 additions & 7 deletions skore/src/skore/persistence/item/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,9 @@ def item_to_object(item: Item) -> Any:
return item.primitive
elif isinstance(item, NumpyArrayItem):
return item.array
elif isinstance(item, PandasDataFrameItem):
elif isinstance(item, PandasDataFrameItem) or isinstance(item, PolarsDataFrameItem):
return item.dataframe
elif isinstance(item, PandasSeriesItem):
return item.series
elif isinstance(item, PolarsDataFrameItem):
return item.dataframe
elif isinstance(item, PolarsSeriesItem):
elif isinstance(item, PandasSeriesItem) or isinstance(item, PolarsSeriesItem):
return item.series
elif isinstance(item, SklearnBaseEstimatorItem):
return item.estimator
Expand All @@ -66,7 +62,7 @@ def item_to_object(item: Item) -> Any:
elif isinstance(item, MediaItem):
return item.media_bytes
elif isinstance(item, PickleItem):
return repr(item.pickle_bytes)
return repr(item.pickle)
else:
raise ValueError(f"Item {item} is not a known item type.")

Expand Down
25 changes: 25 additions & 0 deletions skore/src/skore/persistence/item/pickle_item.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
from functools import cached_property
from pickle import dumps, loads
from typing import Any

from skore.item.item import Item


class PickleItem(Item):
def __init__(
self,
pickle_bytes: bytes,
created_at: str | None = None,
updated_at: str | None = None,
):
super().__init__(created_at, updated_at)

self.pickle_bytes = pickle_bytes

@cached_property
def pickle(self) -> Any:
return loads(self.pickle_bytes)

@classmethod
def factory(cls, object: Any) -> PickleItem:
return cls(dumps(object)

0 comments on commit 06ff76f

Please # to comment.