pull in additional changes from 2.0 branch.
[monitor.git] / monitor / Rpyc / Utils.py
diff --git a/monitor/Rpyc/Utils.py b/monitor/Rpyc/Utils.py
new file mode 100644 (file)
index 0000000..00963c2
--- /dev/null
@@ -0,0 +1,265 @@
+"""\r
+convenience utilities:\r
+ * provides dir(), getattr(), hasattr(), help() and reload() that support netproxies\r
+ * provides obtain() for really transfering remote objects\r
+ * provides upload() and download() for working with files\r
+ * provides a nice interface for remote shell operations (like os.system)\r
+   and a openning a remote python interpreter (for debugging, etc.)\r
+\r
+i removed lots of stuff from the __all__, keeping only the useful things, \r
+so that import * wouldnt mess your namespace too much\r
+"""\r
+import __builtin__\r
+import sys\r
+import os\r
+import inspect\r
+from NetProxy import NetProxy\r
+\r
+__all__ = [\r
+    "dir", "getattr", "hasattr", "help", "reload", "obtain",\r
+    "upload", "download",\r
+]\r
+\r
+CHUNK_SIZE = 4096\r
+\r
+class UtilityError(Exception): \r
+    pass\r
+\r
+#\r
+# working with netproxies\r
+#\r
+def dir(*obj):\r
+    """a version of dir() that supports NetProxies"""\r
+    if not obj:\r
+        inspect_stack = [inspect.stack()[1][0].f_locals.keys()]\r
+        inspect_stack.sort()\r
+        return inspect_stack\r
+    if not len(obj) == 1:\r
+        raise TypeError("dir expected at most 1 arguments, got %d" % (len(obj),))\r
+    obj = obj[0]\r
+    if isinstance(obj, NetProxy):\r
+        return obj.__dict__["____conn"].modules["__builtin__"].dir(obj)\r
+    else:\r
+        return __builtin__.dir(obj)\r
+\r
+def getattr(obj, name, *default):\r
+    """a version of getattr() that supports NetProxies"""\r
+    if len(default) > 1:\r
+        raise TypeError("getattr expected at most 3 arguments, got %d" % (2 + len(default),))\r
+    if isinstance(obj, NetProxy):\r
+        try:\r
+            return obj.__getattr__(name)\r
+        except AttributeError:\r
+            if not default:\r
+                raise\r
+            return default[0]\r
+    else:\r
+        return __builtin__.getattr(obj, name, *default)\r
+\r
+def hasattr(obj, name):\r
+    """a version of hasattr() that supports NetProxies"""\r
+    try:\r
+        getattr(obj, name)\r
+    except AttributeError:\r
+        return False\r
+    else:\r
+        return True\r
+\r
+class _Helper(object):\r
+    """a version of help() that supports NetProxies"""\r
+    def __repr__(self):\r
+        return repr(__builtin__.help)\r
+    def __call__(self, obj = None):\r
+        if isinstance(obj, NetProxy):\r
+            print "Help on NetProxy object for an instance of %r:" % (obj.__getattr__("__class__").__name__,)\r
+            print\r
+            print "Doc:"\r
+            print obj.__doc__\r
+            print\r
+            print "Members:"\r
+            print dir(obj)\r
+        else:\r
+            __builtin__.help(obj)\r
+help = _Helper()\r
+\r
+def reload(module):\r
+    """a version of reload() that supports NetProxies"""\r
+    if isinstance(module, NetProxy):\r
+        return module.__dict__["____conn"].modules["__builtin__"].reload(module)\r
+    else:\r
+        return __builtin__.reload(module)\r
+\r
+def obtain(proxy):\r
+    """transfers a remote object to this process. this is done by pickling it, so it\r
+    must be a picklable object, and should be immutable (otherwise the local object\r
+    may be different from the remote one, which may cause problems). generally speaking, \r
+    you should not obtain remote objects, as NetProxies provide a stronger mechanism.\r
+    but there are times when you want to get the real object in your hand, for pickling\r
+    it locally (e.g., storing test results in a file), or if the connection is too slow."""\r
+    import cPickle\r
+    dumped = proxy.__dict__["____conn"].modules.cPickle.dumps(proxy)\r
+    return cPickle.loads(dumped)\r
+\r
+def getconn(obj):\r
+    """returns the connection of a NetProxy"""\r
+    if "____conn" not in obj.__dict__:\r
+        raise TypeError("`obj` is not a NetProxy")\r
+    return proxy.__dict__["____conn"]\r
+\r
+#\r
+# working with files\r
+#\r
+def upload(conn, localpath, remotepath, *a, **k):\r
+    """uploads a file or a directory recursively (depending on what `localpath` is)"""\r
+    if os.path.isdir(localpath):\r
+        upload_dir(conn, localpath, remotepath, *a, **k)\r
+    elif os.path.isfile(localpath):\r
+        upload_file(conn, localpath, remotepath, *a, **k)\r
+    else:\r
+        raise UtilityError("can only upload files or directories")\r
+\r
+def download(conn, remotepath, localpath, *a, **k):\r
+    """downloads a file or a directory recursively (depending on what `remotepath` is)"""\r
+    if conn.modules.os.path.isdir(remotepath):\r
+        download_dir(conn, remotepath, localpath, *a, **k)\r
+    elif conn.modules.os.path.isfile(remotepath):\r
+        download_file(conn, remotepath, localpath, *a, **k)\r
+    else:\r
+        raise UtilityError("can only download files or directories")\r
+\r
+def upload_file(conn, localpath, remotepath):\r
+    lf = open(localpath, "rb")\r
+    rf = conn.modules.__builtin__.open(remotepath, "wb")\r
+    while True:\r
+        chunk = lf.read(CHUNK_SIZE)\r
+        if not chunk:\r
+            break\r
+        rf.write(chunk)\r
+    lf.close()\r
+    rf.close()\r
+\r
+def download_file(conn, remotepath, localpath):\r
+    lf = open(localpath, "wb")\r
+    rf = conn.modules.__builtin__.open(remotepath, "rb")\r
+    while True:\r
+        chunk = rf.read(CHUNK_SIZE)\r
+        if not chunk:\r
+            break\r
+        lf.write(chunk)\r
+    lf.close()\r
+    rf.close()\r
+    \r
+def upload_dir(conn, localpath, remotepath, extensions = [""]):\r
+    if not conn.modules.os.path.exists(remotepath):\r
+        conn.modules.os.makedirs(remotepath)\r
+    for fn in os.listdir(localpath):\r
+        lfn = os.path.join(localpath, fn)\r
+        rfn = conn.modules.os.path.join(remotepath, fn)\r
+        if os.path.isdir(lfn):\r
+            upload_dir(conn, lfn, rfn, extensions)\r
+        elif os.path.isfile(lfn):\r
+            for ext in extensions:\r
+                if fn.endswith(ext):\r
+                    upload_file(conn, lfn, rfn)\r
+                    break\r
+\r
+def download_dir(conn, remotepath, localpath, extensions = [""]):\r
+    if not os.path.exists(localpath):\r
+        os.makedirs(localpath)\r
+    for fn in conn.modules.os.listdir(remotepath):\r
+        lfn = os.path.join(localpath, fn)\r
+        rfn = conn.modules.os.path.join(remotepath, fn)\r
+        if conn.modules.os.path.isdir(lfn):\r
+            download_dir(conn, rfn, lfn, extensions)\r
+        elif conn.modules.os.path.isfile(lfn):\r
+            for ext in extensions:\r
+                if fn.endswith(ext):\r
+                    download_file(conn, rfn, lfn)\r
+                    break\r
+\r
+#\r
+# distributing modules between hosts\r
+#\r
+def upload_package(conn, module, remotepath = None):\r
+    """uploads the given package to the server, storing it in `remotepath`. if \r
+    remotepath is None, it defaults to the server's site-packages. if the package\r
+    already exists, it is overwritten.\r
+    usage:\r
+        import xml\r
+        upload_package(conn, xml)"""\r
+    if remotepath is None:\r
+        remotepath = conn.modules["distutils.sysconfig"].get_python_lib()\r
+    localpath = os.path.dirname(module.__file__)\r
+    upload_dir(conn, localpath, remotepath, [".py", ".pyd", ".dll", ".so", ".zip"])\r
+\r
+def update_module(conn, module):\r
+    """updates an existing module on the server. the local module is transfered to the\r
+    server, overwriting the old one, and is reloaded. \r
+    usage:\r
+        import xml.dom.minidom\r
+        upload_module(conn, xml.dom.minidom)"""\r
+    remote_module = conn.modules[module.__name__]\r
+    local_file = inspect.getsourcefile(module)\r
+    remote_file = inspect.getsourcefile(remote_module)\r
+    upload_file(conn, local_filem, remote_file)\r
+    reload(remote_module)\r
+\r
+#\r
+# remote shell and interpreter\r
+#\r
+def _redirect_std(conn):\r
+    rsys = conn.modules.sys\r
+    orig = (rsys.stdin, rsys.stdout, rsys.stderr)\r
+    rsys.stdin = sys.stdin\r
+    rsys.stdout = sys.stdout\r
+    rsys.stderr = sys.stderr\r
+    return orig\r
+\r
+def _restore_std(conn, (stdin, stdout, stderr)):\r
+    rsys = conn.modules.sys\r
+    rsys.stdin = stdin\r
+    rsys.stdout = stdout\r
+    rsys.stderr = stderr\r
+    \r
+def remote_shell(conn, command = None):\r
+    """runs the given command on the server, just like os.system, but takes care\r
+    of redirecting the server's stdout/stdin to the client"""\r
+    # BUG: on windows, there's a problem with redirecting the output of spawned command.\r
+    # it runs fine and all, just the client can't see the output. again, windows sucks.\r
+    if command is None:\r
+        if sys.platform == "win32":\r
+            command = "%COMSPEC%"\r
+        else:\r
+            command = "/bin/sh"\r
+    try:\r
+        orig = _redirect_std(conn)\r
+        return conn.modules.os.system(command)\r
+    finally:\r
+        _restore_std(conn, orig)\r
+    \r
+def remote_interpreter(conn, namespace = None):\r
+    """starts an interactive interpreter on the server"""\r
+    if namespace is None:\r
+        #namespace = inspect.stack()[1][0].f_globals.copy()\r
+        #namespace.update(inspect.stack()[1][0].f_locals)\r
+        namespace = {"conn" : conn}\r
+    try:\r
+        orig = _redirect_std(conn)\r
+        return conn.modules["Rpyc.Utils"]._remote_interpreter_server_side(**namespace)\r
+    finally:\r
+        _restore_std(conn, orig)\r
+\r
+def _remote_interpreter_server_side(**namespace):\r
+    import code\r
+    namespace.update(globals())\r
+    code.interact(local = namespace)\r
+\r
+def remote_post_mortem(conn):\r
+    """a version of pdb.pm() that operates on exceptions at the remote side of the connection"""\r
+    import pdb\r
+    pdb.post_mortem(c.modules.sys.last_traceback)\r
+\r
+\r
+\r
+\r
+\r