add possibility to define connect and disconnect callbacks. Further, …

…reset stdin, stdout and stderr in case connection breaks

RPyCRobotRemote/RPyCRobotRemoteServer.py

@@ -12,13 +12,15 @@
 # pylint: disable=E0611
 from rpyc.lib.compat import execute
 # pylint: enable=E0611
+from rpyc.utils.helpers import classpartial
 from rpyc.utils.server import ThreadedServer
 
 
 LOGGER = logging.getLogger('RPyCRobotRemote.Server')
 LOGGER.setLevel(logging.INFO)
 del LOGGER
 
+UNDEFINED = object()
 
 class RPyCRobotRemoteServer:
     """
@@ -63,13 +65,35 @@ def __init__(self, library):
                 super().__init__()
                 self.namespace = {}
                 self._library = library
+                self._stdin = UNDEFINED
+                self._stdout = UNDEFINED
+                self._stderr = UNDEFINED
 
             if allow_remote_stop:
                 @staticmethod
                 def stop_remote_server():
                     """stop server from remote"""
                     self.stop()
 
+            def on_connect(self, conn):
+                on_connect = getattr(self._library, '_on_connect', None)
+                if on_connect:
+                    on_connect()
+
+            def on_disconnect(self, conn):
+                if sys.stdin is self._stdin:
+                    sys.stdin = sys.__stdin__
+
+                if sys.stdout is self._stdout:
+                    sys.stdout = sys.__stdout__
+
+                if sys.stderr is self._stderr:
+                    sys.stderr = sys.__stderr__
+
+                on_disconnect = getattr(self._library, '_on_disconnect', None)
+                if on_disconnect:
+                    on_disconnect()
+
             def execute(self, text):
                 """execute arbitrary code (using ``exec``)"""
                 execute(text, self.namespace)
@@ -113,6 +137,7 @@ def stdin(self):
 
             @stdin.setter
             def stdin(self, value: TextIO):
+                self._stdin = value
                 sys.stdin = value
 
             @property
@@ -122,6 +147,7 @@ def stdout(self):
 
             @stdout.setter
             def stdout(self, value: TextIO):
+                self._stdout = value
                 sys.stdout = value
 
             @property
@@ -131,6 +157,7 @@ def stderr(self):
 
             @stderr.setter
             def stderr(self, value: TextIO):
+                self._stderr = value
                 sys.stderr = value
 
             def _rpyc_setattr(self, name: str, value):
@@ -173,8 +200,9 @@ def _rpyc_setattr(self, name: str, value):
             )
         # pylint: enable=duplicate-code
 
+        service = classpartial(Service, library)
         self._server = server(
-            Service(library),
+            service,
             hostname=host,
             port=port,
             ipv6=ipv6,