Rpyc/Utils.py

   1 """\r
   2 convenience utilities:\r
   3  * provides dir(), getattr(), hasattr(), help() and reload() that support netproxies\r
   4  * provides obtain() for really transfering remote objects\r
   5  * provides upload() and download() for working with files\r
   6  * provides a nice interface for remote shell operations (like os.system)\r
   7    and a openning a remote python interpreter (for debugging, etc.)\r
   8 \r
   9 i removed lots of stuff from the __all__, keeping only the useful things, \r
  10 so that import * wouldnt mess your namespace too much\r
  11 """\r
  12 import __builtin__\r
  13 import sys\r
  14 import os\r
  15 import inspect\r
  16 from NetProxy import NetProxy\r
  17 \r
  18 __all__ = [\r
  19     "dir", "getattr", "hasattr", "help", "reload", "obtain",\r
  20     "upload", "download",\r
  21 ]\r
  22 \r
  23 CHUNK_SIZE = 4096\r
  24 \r
  25 class UtilityError(Exception): \r
  26     pass\r
  27 \r
  28 #\r
  29 # working with netproxies\r
  30 #\r
  31 def dir(*obj):\r
  32     """a version of dir() that supports NetProxies"""\r
  33     if not obj:\r
  34         inspect_stack = [inspect.stack()[1][0].f_locals.keys()]\r
  35         inspect_stack.sort()\r
  36         return inspect_stack\r
  37     if not len(obj) == 1:\r
  38         raise TypeError("dir expected at most 1 arguments, got %d" % (len(obj),))\r
  39     obj = obj[0]\r
  40     if isinstance(obj, NetProxy):\r
  41         return obj.__dict__["____conn"].modules["__builtin__"].dir(obj)\r
  42     else:\r
  43         return __builtin__.dir(obj)\r
  44 \r
  45 def getattr(obj, name, *default):\r
  46     """a version of getattr() that supports NetProxies"""\r
  47     if len(default) > 1:\r
  48         raise TypeError("getattr expected at most 3 arguments, got %d" % (2 + len(default),))\r
  49     if isinstance(obj, NetProxy):\r
  50         try:\r
  51             return obj.__getattr__(name)\r
  52         except AttributeError:\r
  53             if not default:\r
  54                 raise\r
  55             return default[0]\r
  56     else:\r
  57         return __builtin__.getattr(obj, name, *default)\r
  58 \r
  59 def hasattr(obj, name):\r
  60     """a version of hasattr() that supports NetProxies"""\r
  61     try:\r
  62         getattr(obj, name)\r
  63     except AttributeError:\r
  64         return False\r
  65     else:\r
  66         return True\r
  67 \r
  68 class _Helper(object):\r
  69     """a version of help() that supports NetProxies"""\r
  70     def __repr__(self):\r
  71         return repr(__builtin__.help)\r
  72     def __call__(self, obj = None):\r
  73         if isinstance(obj, NetProxy):\r
  74             print "Help on NetProxy object for an instance of %r:" % (obj.__getattr__("__class__").__name__,)\r
  75             print\r
  76             print "Doc:"\r
  77             print obj.__doc__\r
  78             print\r
  79             print "Members:"\r
  80             print dir(obj)\r
  81         else:\r
  82             __builtin__.help(obj)\r
  83 help = _Helper()\r
  84 \r
  85 def reload(module):\r
  86     """a version of reload() that supports NetProxies"""\r
  87     if isinstance(module, NetProxy):\r
  88         return module.__dict__["____conn"].modules["__builtin__"].reload(module)\r
  89     else:\r
  90         return __builtin__.reload(module)\r
  91 \r
  92 def obtain(proxy):\r
  93     """transfers a remote object to this process. this is done by pickling it, so it\r
  94     must be a picklable object, and should be immutable (otherwise the local object\r
  95     may be different from the remote one, which may cause problems). generally speaking, \r
  96     you should not obtain remote objects, as NetProxies provide a stronger mechanism.\r
  97     but there are times when you want to get the real object in your hand, for pickling\r
  98     it locally (e.g., storing test results in a file), or if the connection is too slow."""\r
  99     import cPickle\r
 100     dumped = proxy.__dict__["____conn"].modules.cPickle.dumps(proxy)\r
 101     return cPickle.loads(dumped)\r
 102 \r
 103 def getconn(obj):\r
 104     """returns the connection of a NetProxy"""\r
 105     if "____conn" not in obj.__dict__:\r
 106         raise TypeError("`obj` is not a NetProxy")\r
 107     return proxy.__dict__["____conn"]\r
 108 \r
 109 #\r
 110 # working with files\r
 111 #\r
 112 def upload(conn, localpath, remotepath, *a, **k):\r
 113     """uploads a file or a directory recursively (depending on what `localpath` is)"""\r
 114     if os.path.isdir(localpath):\r
 115         upload_dir(conn, localpath, remotepath, *a, **k)\r
 116     elif os.path.isfile(localpath):\r
 117         upload_file(conn, localpath, remotepath, *a, **k)\r
 118     else:\r
 119         raise UtilityError("can only upload files or directories")\r
 120 \r
 121 def download(conn, remotepath, localpath, *a, **k):\r
 122     """downloads a file or a directory recursively (depending on what `remotepath` is)"""\r
 123     if conn.modules.os.path.isdir(remotepath):\r
 124         download_dir(conn, remotepath, localpath, *a, **k)\r
 125     elif conn.modules.os.path.isfile(remotepath):\r
 126         download_file(conn, remotepath, localpath, *a, **k)\r
 127     else:\r
 128         raise UtilityError("can only download files or directories")\r
 129 \r
 130 def upload_file(conn, localpath, remotepath):\r
 131     lf = open(localpath, "rb")\r
 132     rf = conn.modules.__builtin__.open(remotepath, "wb")\r
 133     while True:\r
 134         chunk = lf.read(CHUNK_SIZE)\r
 135         if not chunk:\r
 136             break\r
 137         rf.write(chunk)\r
 138     lf.close()\r
 139     rf.close()\r
 140 \r
 141 def download_file(conn, remotepath, localpath):\r
 142     lf = open(localpath, "wb")\r
 143     rf = conn.modules.__builtin__.open(remotepath, "rb")\r
 144     while True:\r
 145         chunk = rf.read(CHUNK_SIZE)\r
 146         if not chunk:\r
 147             break\r
 148         lf.write(chunk)\r
 149     lf.close()\r
 150     rf.close()\r
 151     \r
 152 def upload_dir(conn, localpath, remotepath, extensions = [""]):\r
 153     if not conn.modules.os.path.exists(remotepath):\r
 154         conn.modules.os.makedirs(remotepath)\r
 155     for fn in os.listdir(localpath):\r
 156         lfn = os.path.join(localpath, fn)\r
 157         rfn = conn.modules.os.path.join(remotepath, fn)\r
 158         if os.path.isdir(lfn):\r
 159             upload_dir(conn, lfn, rfn, extensions)\r
 160         elif os.path.isfile(lfn):\r
 161             for ext in extensions:\r
 162                 if fn.endswith(ext):\r
 163                     upload_file(conn, lfn, rfn)\r
 164                     break\r
 165 \r
 166 def download_dir(conn, remotepath, localpath, extensions = [""]):\r
 167     if not os.path.exists(localpath):\r
 168         os.makedirs(localpath)\r
 169     for fn in conn.modules.os.listdir(remotepath):\r
 170         lfn = os.path.join(localpath, fn)\r
 171         rfn = conn.modules.os.path.join(remotepath, fn)\r
 172         if conn.modules.os.path.isdir(lfn):\r
 173             download_dir(conn, rfn, lfn, extensions)\r
 174         elif conn.modules.os.path.isfile(lfn):\r
 175             for ext in extensions:\r
 176                 if fn.endswith(ext):\r
 177                     download_file(conn, rfn, lfn)\r
 178                     break\r
 179 \r
 180 #\r
 181 # distributing modules between hosts\r
 182 #\r
 183 def upload_package(conn, module, remotepath = None):\r
 184     """uploads the given package to the server, storing it in `remotepath`. if \r
 185     remotepath is None, it defaults to the server's site-packages. if the package\r
 186     already exists, it is overwritten.\r
 187     usage:\r
 188         import xml\r
 189         upload_package(conn, xml)"""\r
 190     if remotepath is None:\r
 191         remotepath = conn.modules["distutils.sysconfig"].get_python_lib()\r
 192     localpath = os.path.dirname(module.__file__)\r
 193     upload_dir(conn, localpath, remotepath, [".py", ".pyd", ".dll", ".so", ".zip"])\r
 194 \r
 195 def update_module(conn, module):\r
 196     """updates an existing module on the server. the local module is transfered to the\r
 197     server, overwriting the old one, and is reloaded. \r
 198     usage:\r
 199         import xml.dom.minidom\r
 200         upload_module(conn, xml.dom.minidom)"""\r
 201     remote_module = conn.modules[module.__name__]\r
 202     local_file = inspect.getsourcefile(module)\r
 203     remote_file = inspect.getsourcefile(remote_module)\r
 204     upload_file(conn, local_filem, remote_file)\r
 205     reload(remote_module)\r
 206 \r
 207 #\r
 208 # remote shell and interpreter\r
 209 #\r
 210 def _redirect_std(conn):\r
 211     rsys = conn.modules.sys\r
 212     orig = (rsys.stdin, rsys.stdout, rsys.stderr)\r
 213     rsys.stdin = sys.stdin\r
 214     rsys.stdout = sys.stdout\r
 215     rsys.stderr = sys.stderr\r
 216     return orig\r
 217 \r
 218 def _restore_std(conn, (stdin, stdout, stderr)):\r
 219     rsys = conn.modules.sys\r
 220     rsys.stdin = stdin\r
 221     rsys.stdout = stdout\r
 222     rsys.stderr = stderr\r
 223     \r
 224 def remote_shell(conn, command = None):\r
 225     """runs the given command on the server, just like os.system, but takes care\r
 226     of redirecting the server's stdout/stdin to the client"""\r
 227     # BUG: on windows, there's a problem with redirecting the output of spawned command.\r
 228     # it runs fine and all, just the client can't see the output. again, windows sucks.\r
 229     if command is None:\r
 230         if sys.platform == "win32":\r
 231             command = "%COMSPEC%"\r
 232         else:\r
 233             command = "/bin/sh"\r
 234     try:\r
 235         orig = _redirect_std(conn)\r
 236         return conn.modules.os.system(command)\r
 237     finally:\r
 238         _restore_std(conn, orig)\r
 239     \r
 240 def remote_interpreter(conn, namespace = None):\r
 241     """starts an interactive interpreter on the server"""\r
 242     if namespace is None:\r
 243         #namespace = inspect.stack()[1][0].f_globals.copy()\r
 244         #namespace.update(inspect.stack()[1][0].f_locals)\r
 245         namespace = {"conn" : conn}\r
 246     try:\r
 247         orig = _redirect_std(conn)\r
 248         return conn.modules["Rpyc.Utils"]._remote_interpreter_server_side(**namespace)\r
 249     finally:\r
 250         _restore_std(conn, orig)\r
 251 \r
 252 def _remote_interpreter_server_side(**namespace):\r
 253     import code\r
 254     namespace.update(globals())\r
 255     code.interact(local = namespace)\r
 256 \r
 257 def remote_post_mortem(conn):\r
 258     """a version of pdb.pm() that operates on exceptions at the remote side of the connection"""\r
 259     import pdb\r
 260     pdb.post_mortem(c.modules.sys.last_traceback)\r
 261 \r
 262 \r
 263 \r
 264 \r
 265 \r