diff --git a/src/neptune_scale/api/run.py b/src/neptune_scale/api/run.py index f6d05f3..41bfbe4 100644 --- a/src/neptune_scale/api/run.py +++ b/src/neptune_scale/api/run.py @@ -24,10 +24,7 @@ from neptune_api.proto.neptune_pb.ingest.v1.common_pb2 import Run as CreateRun from neptune_api.proto.neptune_pb.ingest.v1.pub.ingest_pb2 import RunOperation -from neptune_scale.api.attribute import ( - Attribute, - AttributeStore, -) +from neptune_scale.api.attribute import AttributeStore from neptune_scale.api.validation import ( verify_dict_type, verify_max_length, @@ -390,12 +387,6 @@ def _create_run( ) self._operations_queue.enqueue(operation=operation) - def __getitem__(self, key: str) -> Attribute: - return self._attr_store[key] - - def __setitem__(self, key: str, value: Any) -> None: - self._attr_store[key] = value - def log_metrics( self, data: dict[str, Union[float, int]], diff --git a/src/neptune_scale/legacy.py b/src/neptune_scale/legacy.py new file mode 100644 index 0000000..3ce25ec --- /dev/null +++ b/src/neptune_scale/legacy.py @@ -0,0 +1,28 @@ +from typing import Any + +import neptune_scale.api.run +from neptune_scale.api.attribute import Attribute + +__all__ = ("Run",) + + +class Run(neptune_scale.api.run.Run): + """This class extends the main Run class with a dict-like API compatible (on a basic level) + with the legacy neptune-client package. + + Example: + + from neptune_scale.legacy import Run + + run = Run(...) + run['foo'] = 1 + run['metrics/loss'].append(0.5, step=10) + + run.close() + """ + + def __getitem__(self, key: str) -> Attribute: + return self._attr_store[key] + + def __setitem__(self, key: str, value: Any) -> None: + self._attr_store[key] = value diff --git a/tests/unit/test_attribute.py b/tests/unit/test_attribute.py index e919dbd..481651c 100644 --- a/tests/unit/test_attribute.py +++ b/tests/unit/test_attribute.py @@ -7,8 +7,8 @@ mark, ) -from neptune_scale import Run from neptune_scale.api.attribute import cleanup_path +from neptune_scale.legacy import Run @fixture