Create proper interface for Twitter::IdentityMap

sferik · Jul 17, 2012 · 8996c37 · 8996c37
1 parent ff9e036
commit 8996c37
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 12 deletions.
diff --git a/lib/twitter/base.rb b/lib/twitter/base.rb
@@ -25,12 +25,9 @@ def self.attr_reader(*attrs)
     # @return [Twitter::Base]
     def self.fetch(attrs)
       return unless Twitter.identity_map
-
-      Twitter.identity_map[self] ||= {}
-      if object = Twitter.identity_map[self][Marshal.dump(attrs)]
+      if object = Twitter.identity_map.fetch(self, Marshal.dump(attrs))
         return object
       end
-
       return yield if block_given?
       raise Twitter::IdentityMapKeyError, 'key not found'
     end
@@ -41,9 +38,7 @@ def self.fetch(attrs)
     # @return [Twitter::Base]
     def self.store(object)
       return object unless Twitter.identity_map
-
-      Twitter.identity_map[self] ||= {}
-      Twitter.identity_map[self][Marshal.dump(object.attrs)] = object
+      Twitter.identity_map.store(Marshal.dump(object.attrs), object)
     end
 
     # Returns a new object based on the response hash

diff --git a/lib/twitter/identity.rb b/lib/twitter/identity.rb
@@ -7,9 +7,8 @@ def self.fetch(attrs)
       return unless Twitter.identity_map
 
       id = attrs[:id]
-      Twitter.identity_map[self] ||= {}
-      if id && Twitter.identity_map[self][id]
-        return Twitter.identity_map[self][id].update(attrs)
+      if id && object = Twitter.identity_map.fetch(self, id)
+        return object.update(attrs)
       end
 
       return yield if block_given?
@@ -21,8 +20,8 @@ def self.fetch(attrs)
     # @param attrs [Hash]
     # @return [Twitter::Identity]
     def self.store(object)
-      Twitter.identity_map[self] ||= {}
-      object.id && Twitter.identity_map[self][object.id] = object || super(object)
+      return object unless Twitter.identity_map
+      Twitter.identity_map.store(object.id, object)
     end
 
     # Initializes a new object

diff --git a/lib/twitter/identity_map.rb b/lib/twitter/identity_map.rb
@@ -3,6 +3,22 @@ module Twitter
   # Tracks objects to help ensure that each object gets loaded only once.
   # See: http://www.martinfowler.com/eaaCatalog/identityMap.html
   class IdentityMap < Hash
+
+    # @param klass
+    # @param key
+    # @return [Object]
+    def fetch(klass, key)
+      self[klass] && self[klass][key]
+    end
+
+    # @param key
+    # @param object
+    # @return [Object]
+    def store(key, object)
+      self[object.class] ||= {}
+      self[object.class][key] = object
+    end
+
   end
 
   # Inherit from KeyError when Ruby 1.8 compatibility is removed