Merge pull request #141 from mrkn/use_weak_map

Use WeakMap for WrapperObjectCache

Author: mrkn

lib/pycall/wrapper_object_cache.rb

@@ -1,5 +1,47 @@
 module PyCall
   class WrapperObjectCache
+
+    begin
+      ObjectSpace::WeakMap.new[42] = Object.new
+    rescue
+      WMAP_SUPPORT_INT_KEY = false
+    else
+      WMAP_SUPPORT_INT_KEY = true
+    end
+
+    if WMAP_SUPPORT_INT_KEY
+      def self.get_key(pyptr)
+        pyptr.__address__
+      end
+    else
+      class Key
+        @address_key_map = {}
+
+        def self.[](address)
+          @address_key_map[address] ||= new(address)
+        end
+
+        def initialize(address)
+          @address = address
+        end
+
+        attr_reader :address
+
+        def ==(other)
+          case other
+          when Key
+            self.address == other.address
+          else
+            super
+          end
+        end
+      end
+
+      def self.get_key(pyptr)
+        Key[pyptr.__address__]
+      end
+    end
+
     def initialize(*restricted_pytypes)
       unless restricted_pytypes.empty?
         restricted_pytypes.each do |pytype|
@@ -8,9 +50,7 @@ def initialize(*restricted_pytypes)
         end
       end
       @restricted_pytypes = restricted_pytypes
-      @wrapper_object_table = {}
-      @wrapped_pyptr_table = {}
-      @weakref_table = {}
+      @wrapper_object_table = ObjectSpace::WeakMap.new
     end
 
     def lookup(pyptr)
@@ -25,16 +65,14 @@ def lookup(pyptr)
         end
       end
 
-      wrapper_object_id = @wrapper_object_table[pyptr.__address__]
-      if wrapper_object_id
-        wrapper_object = ObjectSpace._id2ref(wrapper_object_id) rescue nil
-        return wrapper_object if wrapper_object
+      key = self.class.get_key(pyptr)
+      wrapper_object = @wrapper_object_table[key]
+      unless wrapper_object
+        wrapper_object = yield(pyptr)
+        check_wrapper_object(wrapper_object)
+        @wrapper_object_table[key] = wrapper_object
       end
 
-      wrapper_object = yield(pyptr)
-      check_wrapper_object(wrapper_object)
-      register_wrapper_object(pyptr, wrapper_object)
-
       wrapper_object
     end
 
@@ -43,19 +81,5 @@ def check_wrapper_object(wrapper_object)
         raise TypeError, "unexpected wrapper object (expected an object extended by PyObjectWrapper)"
       end
     end
-
-    def register_wrapper_object(pyptr, wrapper_object)
-      @wrapper_object_table[pyptr.__address__] = wrapper_object.__id__
-      @wrapped_pyptr_table[wrapper_object.__id__] = pyptr.__address__
-      ObjectSpace.define_finalizer(wrapper_object, &method(:unregister_wrapper_object))
-      # TODO: weakref
-      self
-    end
-
-    def unregister_wrapper_object(wrapper_object_id)
-      pyptr_addr = @wrapped_pyptr_table.delete(wrapper_object_id)
-      @wrapper_object_table.delete(pyptr_addr) if pyptr_addr
-      self
-    end
   end
 end