Rpyc/Boxing.py

   1 import sys\r
   2 import traceback\r
   3 import cPickle as pickle\r
   4 from weakref import WeakValueDictionary\r
   5 from Lib import ImmDict\r
   6 from NetProxy import NetProxy, SyncNetProxy\r
   7 \r
   8 \r
   9 class BoxingError(Exception):\r
  10     pass\r
  11 class NestedException(Exception): \r
  12     pass\r
  13 \r
  14 PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL\r
  15 TYPE_SIMPLE = 0\r
  16 TYPE_PROXY = 1\r
  17 TYPE_TUPLE = 2\r
  18 TYPE_SLICE = 3\r
  19 TYPE_LOCAL_PROXY = 4\r
  20 TYPE_IMMDICT = 5\r
  21 simple_types = (\r
  22     bool, \r
  23     int, \r
  24     long, \r
  25     float, \r
  26     complex, \r
  27     basestring, \r
  28     type(None),\r
  29 )\r
  30 \r
  31 def dump_exception(typ, val, tb):\r
  32     """dumps the given exception using pickle (since not all exceptions are picklable)"""\r
  33     tbtext = "".join(traceback.format_exception(typ, val, tb))\r
  34     sys.last_type, sys.last_value, sys.last_traceback = typ, val, tb\r
  35     try:\r
  36         pickled_exc = pickle.dumps((typ, val, tbtext), PICKLE_PROTOCOL)\r
  37     except pickle.PicklingError, ex:\r
  38         newval = NestedException("pickling error %s\nexception type: %r\nexception object: %s" % (ex, typ, val))\r
  39         pickled_exc = pickle.dumps((NestedException, newval, tbtext), PICKLE_PROTOCOL)\r
  40     return pickled_exc\r
  41 \r
  42 def load_exception(package):\r
  43     """returns an exception object"""\r
  44     try:\r
  45         return pickle.loads(package)\r
  46     except pickle.PicklingError, ex:\r
  47         return NestedException("failed to unpickle remote exception -- %r" % (ex,))\r
  48 \r
  49 class Box(object):\r
  50     """a box is where local objects are stored, and remote proxies are created"""\r
  51     def __init__(self, conn):\r
  52         self.conn = conn\r
  53         self.objects = {}\r
  54         self.proxy_cache = WeakValueDictionary()\r
  55 \r
  56     def close(self):\r
  57         del self.conn\r
  58         del self.objects\r
  59         del self.proxy_cache\r
  60     \r
  61     def __getitem__(self, oid):\r
  62         return self.objects[oid][1]\r
  63 \r
  64     def _box(self, obj):\r
  65         if isinstance(obj, simple_types):\r
  66             return TYPE_SIMPLE, obj\r
  67         elif isinstance(obj, slice):\r
  68             return TYPE_SLICE, (obj.start, obj.stop, obj.step)\r
  69         elif isinstance(obj, NetProxy) and obj.__dict__["____conn"] is self.conn:\r
  70             return TYPE_LOCAL_PROXY, obj.__dict__["____oid"]\r
  71         elif isinstance(obj, tuple):\r
  72             if obj:\r
  73                 return TYPE_TUPLE, [self._box(subobj) for subobj in obj]\r
  74             else:\r
  75                 return TYPE_SIMPLE, ()\r
  76         elif isinstance(obj, ImmDict):\r
  77             if not obj.dict:\r
  78                 return TYPE_SIMPLE, {}\r
  79             else:\r
  80                 return TYPE_IMMDICT, [(self._box(k), self._box(v)) for k, v in obj.items()]\r
  81         else:\r
  82             oid = id(obj)\r
  83             if oid not in self.objects:\r
  84                 self.objects[oid] = [0, obj]\r
  85             return TYPE_PROXY, oid\r
  86         \r
  87     def _unbox(self, (type, value)):\r
  88         if type == TYPE_SIMPLE:\r
  89             return value\r
  90         elif type == TYPE_TUPLE:\r
  91             return tuple([self._unbox(subobj) for subobj in value])\r
  92         elif type == TYPE_SLICE:\r
  93             return slice(*value)\r
  94         elif type == TYPE_LOCAL_PROXY:\r
  95             return self[value]\r
  96         elif type == TYPE_IMMDICT:\r
  97             return dict([(self._unbox(k), self._unbox(v)) for k, v in value])\r
  98         elif type == TYPE_PROXY:\r
  99             if value in self.proxy_cache:\r
 100                 proxy = self.proxy_cache[value]\r
 101             else:\r
 102                 proxy = SyncNetProxy(self.conn, value)\r
 103                 self.proxy_cache[value] = proxy\r
 104             return proxy\r
 105         else:\r
 106             raise BoxingError("invalid boxed object type", type, value)\r
 107         \r
 108     def incref(self, oid):\r
 109         self.objects[oid][0] += 1\r
 110 \r
 111     def decref(self, oid):\r
 112         self.objects[oid][0] -= 1\r
 113         if self.objects[oid][0] <= 0:\r
 114             del self.objects[oid]\r
 115     \r
 116     def pack(self, obj):\r
 117         """packs an object (returns a package)"""\r
 118         return pickle.dumps(self._box(obj), PICKLE_PROTOCOL)\r
 119 \r
 120     def unpack(self, package):\r
 121         """unpacks a package (returns an object)"""\r
 122         return self._unbox(pickle.loads(package))\r
 123 \r
 124 \r