lib/pycall/wrapper_object_cache.rb in pycall-1.3.1 vs lib/pycall/wrapper_object_cache.rb in pycall-1.4.0

@@ -1,18 +1,60 @@
 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)
+          # An instance of Key created here is parmanently cached in @address_key_map.
+          # This behavior is intentional.
+          @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|
           next if pytype.kind_of? PyTypePtr
           raise TypeError, "unexpected type of object in the arguments (#{pytype.class} for PyCall::PyTypePtr)"
         end
       end
       @restricted_pytypes = restricted_pytypes
-      @wrapper_object_table = {}
-      @wrapped_pyptr_table = {}
-      @weakref_table = {}
+      @wrapper_object_table = ObjectSpace::WeakMap.new
     end
 
     def lookup(pyptr)
       # TODO: check pytypeptr type
       unless pyptr.kind_of? PyPtr
@@ -23,39 +65,23 @@
         unless @restricted_pytypes.any? {|pytype| pyptr.kind_of? pytype }
           raise TypeError, "unexpected argument Python type #{pyptr.__ob_type__.__name__} (expected either of them in [#{@restricted_pytypes.map(&:__tp_name__).join(', ')}])"
         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
 
     def check_wrapper_object(wrapper_object)
       unless wrapper_object.kind_of?(PyObjectWrapper)
         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