First approach

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,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)

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

src/__init__.py → 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

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: