Merge pull request #3 from neptune-ai/rj/experimental-multiple-features

Support for multiple features

CHANGELOG.md

@@ -5,3 +5,4 @@
 
 ### Changes
 - Allow to disable handling of remote signals ([#2](https://github.com/neptune-ai/neptune-client-experimental/pull/2))
+- Added support for multiple features ([#3](https://github.com/neptune-ai/neptune-client-experimental/pull/3))

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"

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

src/neptune_experimental/another_feature.py

@@ -0,0 +1,39 @@
+#
+# 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.
+#
+"""
+This is just a boilerplate code to show how to use the override function.
+"""
+__all__ = ["initialize"]
+
+
+from typing import (
+    Any,
+    Callable,
+)
+
+from neptune import Run
+
+from neptune_experimental.utils import override
+
+
+def initialize() -> None:
+    # Monkey patching
+    override(obj=Run, method="__init__", target=init_with_print)
+
+
+def init_with_print(self: "Run", *args: Any, original: Callable[..., Any], **kwargs: Any) -> None:
+    print("That's another feature")
+    original(self, *args, **kwargs)

src/neptune_experimental/remote_signals.py

@@ -0,0 +1,62 @@
+#
+# 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,
+    Callable,
+)
+
+from neptune import Run
+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
+    override(obj=Run, method="__init__", target=init_with_enable_remote_signals)
+    override(obj=Run, method="_prepare_background_jobs", target=prepare_background_jobs)
+
+
+def init_with_enable_remote_signals(self: "Run", *args: Any, original: Callable[..., 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
+
+    original(self, *args, **kwargs)
+
+
+def prepare_background_jobs(self: "Run", original: Callable[..., Any]) -> 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

src/__init__.py → src/neptune_experimental/utils.py

@@ -13,3 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+__all__ = ["override"]
+
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+)
+
+
+def override(*, obj: Any, method: str, target: Callable[..., Any]) -> None:
+    source = getattr(obj, method)
+
+    @wraps(source)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        return target(*args, original=source, **kwargs)
+
+    setattr(obj, method, wrapper)

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: