From c415aacb787e2a44a54fcaf34ad59ee08b15a217 Mon Sep 17 00:00:00 2001 From: Rafal Jankowski Date: Thu, 2 Nov 2023 17:47:45 +0100 Subject: [PATCH] First approach --- pyproject.toml | 3 +- src/neptune_experimental/__init__.py | 8 --- src/neptune_experimental/another_feature.py | 30 ++++++++++ src/neptune_experimental/remote_signals.py | 60 +++++++++++++++++++ src/neptune_experimental/run.py | 51 ---------------- .../utils.py} | 16 +++++ tests/unit/test_custom_run.py | 7 --- 7 files changed, 108 insertions(+), 67 deletions(-) create mode 100644 src/neptune_experimental/another_feature.py create mode 100644 src/neptune_experimental/remote_signals.py delete mode 100644 src/neptune_experimental/run.py rename src/{__init__.py => neptune_experimental/utils.py} (66%) diff --git a/pyproject.toml b/pyproject.toml index 68b974e..7534260 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,8 @@ packages = [ ] [tool.poetry.plugins."neptune.extensions"] -"experimental" = "neptune_experimental:initialize" +"experimental_remote_signals" = "neptune_experimental.remote_signals:initialize" +"experimental_another_feature" = "neptune_experimental.another_feature:initialize" [tool.poetry.urls] "Tracker" = "https://github.com/neptune-ai/neptune-client-experimental/issues" diff --git a/src/neptune_experimental/__init__.py b/src/neptune_experimental/__init__.py index 694f59d..8d06af5 100644 --- a/src/neptune_experimental/__init__.py +++ b/src/neptune_experimental/__init__.py @@ -13,11 +13,3 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import neptune - -from neptune_experimental.run import CustomRun - - -def initialize() -> None: - # Monkey patching - neptune.Run = CustomRun diff --git a/src/neptune_experimental/another_feature.py b/src/neptune_experimental/another_feature.py new file mode 100644 index 0000000..de7d3b1 --- /dev/null +++ b/src/neptune_experimental/another_feature.py @@ -0,0 +1,30 @@ +# +# Copyright (c) 2023, Neptune Labs Sp. z o.o. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +__all__ = ["initialize"] + +import neptune + +from neptune_experimental.utils import override + + +def initialize() -> None: + # Monkey patching + neptune.Run.__init__ = override(target=neptune.Run.__init__)(init_with_print) + + +def init_with_print(self, *args, original, **kwargs) -> None: + print("That's another feature") + original(self, *args, **kwargs) diff --git a/src/neptune_experimental/remote_signals.py b/src/neptune_experimental/remote_signals.py new file mode 100644 index 0000000..961d135 --- /dev/null +++ b/src/neptune_experimental/remote_signals.py @@ -0,0 +1,60 @@ +# +# Copyright (c) 2023, Neptune Labs Sp. z o.o. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +__all__ = ["initialize"] + +from typing import Any + +import neptune +from neptune.internal.backgroud_job_list import BackgroundJobList +from neptune.internal.utils import verify_type +from neptune.internal.websockets.websocket_signals_background_job import WebsocketSignalsBackgroundJob + +from neptune_experimental.utils import override + + +def initialize() -> None: + # Monkey patching + neptune.Run.__init__ = override(target=neptune.Run.__init__)(init_with_enable_remote_signals) + neptune.Run._prepare_background_jobs = override(target=neptune.Run._prepare_background_jobs)(prepare_background_jobs) + + +def init_with_enable_remote_signals(self, *args: Any, original, **kwargs: Any) -> None: + print('Hello from "remote_signals" feature!') + enable_remote_signals = kwargs.pop("enable_remote_signals", None) + + if enable_remote_signals is None: + self._enable_remote_signals = True # user did not pass this param in kwargs -> default value + else: + + verify_type("enable_remote_signals", enable_remote_signals, bool) + self._enable_remote_signals = enable_remote_signals + + original(self, *args, **kwargs) + + +def prepare_background_jobs(self, original) -> BackgroundJobList: + background_jobs = original(self) + + if not self._enable_remote_signals: + # filter-out websocket job + background_jobs._jobs = list( + filter( + lambda x: not isinstance(x, WebsocketSignalsBackgroundJob), + background_jobs._jobs, + ) + ) + + return background_jobs diff --git a/src/neptune_experimental/run.py b/src/neptune_experimental/run.py deleted file mode 100644 index 5d52f35..0000000 --- a/src/neptune_experimental/run.py +++ /dev/null @@ -1,51 +0,0 @@ -# -# Copyright (c) 2023, Neptune Labs Sp. z o.o. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -__all__ = ["CustomRun"] - -from typing import Any - -import neptune -from neptune.internal.backgroud_job_list import BackgroundJobList -from neptune.internal.utils import verify_type -from neptune.internal.websockets.websocket_signals_background_job import WebsocketSignalsBackgroundJob - - -class CustomRun(neptune.Run): - def __init__(self, *args: Any, **kwargs: Any) -> None: - enable_remote_signals = kwargs.pop("enable_remote_signals", None) - - if enable_remote_signals is None: - self._enable_remote_signals = True # user did not pass this param in kwargs -> default value - else: - - verify_type("enable_remote_signals", enable_remote_signals, bool) - self._enable_remote_signals = enable_remote_signals - - super().__init__(*args, **kwargs) - - def _prepare_background_jobs(self) -> BackgroundJobList: - background_jobs = super()._prepare_background_jobs() - - if not self._enable_remote_signals: - # filter-out websocket job - background_jobs._jobs = list( - filter( - lambda x: not isinstance(x, WebsocketSignalsBackgroundJob), - background_jobs._jobs, - ) - ) - - return background_jobs diff --git a/src/__init__.py b/src/neptune_experimental/utils.py similarity index 66% rename from src/__init__.py rename to src/neptune_experimental/utils.py index 8d06af5..b245ffe 100644 --- a/src/__init__.py +++ b/src/neptune_experimental/utils.py @@ -13,3 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. # +__all__ = ["override"] + +from functools import wraps + + +def override(target): + print("override called") + + def simple_decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + func(*args, original=target, **kwargs) + + return wrapper + + return simple_decorator diff --git a/tests/unit/test_custom_run.py b/tests/unit/test_custom_run.py index 6374872..075b6fc 100644 --- a/tests/unit/test_custom_run.py +++ b/tests/unit/test_custom_run.py @@ -17,13 +17,6 @@ from neptune import Run from neptune.internal.websockets.websocket_signals_background_job import WebsocketSignalsBackgroundJob -from neptune_experimental.run import CustomRun - - -def test_custom_run(): - with Run(mode="debug") as run: - assert isinstance(run, CustomRun) - def test_disabled_remote_signals(): with Run(mode="debug", enable_remote_signals=False) as run: