From 11d96a71bd64259c0b41e6cc133bf3f0d6d14fc2 Mon Sep 17 00:00:00 2001
From: Rafal Jankowski <rafal.jankowski@neptune.ai>
Date: Thu, 2 Nov 2023 18:32:57 +0100
Subject: [PATCH] Refactor

---
 src/neptune_experimental/another_feature.py |  4 ++--
 src/neptune_experimental/remote_signals.py  |  7 +++----
 src/neptune_experimental/utils.py           | 14 ++++++--------
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/neptune_experimental/another_feature.py b/src/neptune_experimental/another_feature.py
index de7d3b1..e2816e3 100644
--- a/src/neptune_experimental/another_feature.py
+++ b/src/neptune_experimental/another_feature.py
@@ -15,14 +15,14 @@
 #
 __all__ = ["initialize"]
 
-import neptune
+from neptune import Run
 
 from neptune_experimental.utils import override
 
 
 def initialize() -> None:
     # Monkey patching
-    neptune.Run.__init__ = override(target=neptune.Run.__init__)(init_with_print)
+    override(obj=Run, method="__init__", target=init_with_print)
 
 
 def init_with_print(self, *args, original, **kwargs) -> None:
diff --git a/src/neptune_experimental/remote_signals.py b/src/neptune_experimental/remote_signals.py
index 961d135..6f99c1a 100644
--- a/src/neptune_experimental/remote_signals.py
+++ b/src/neptune_experimental/remote_signals.py
@@ -17,7 +17,7 @@
 
 from typing import Any
 
-import neptune
+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
@@ -27,12 +27,11 @@
 
 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)
+    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, *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:
diff --git a/src/neptune_experimental/utils.py b/src/neptune_experimental/utils.py
index b245ffe..bb845e5 100644
--- a/src/neptune_experimental/utils.py
+++ b/src/neptune_experimental/utils.py
@@ -18,14 +18,12 @@
 from functools import wraps
 
 
-def override(target):
-    print("override called")
+def override(*, obj, method, target):
+    old_method = getattr(obj, method)
 
-    def simple_decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            func(*args, original=target, **kwargs)
+    @wraps(old_method)
+    def wrapper(*args, **kwargs):
+        return target(*args, original=old_method, **kwargs)
 
-        return wrapper
+    setattr(obj, method, wrapper)
 
-    return simple_decorator