lib/mem.rb in mem-0.1.4 vs lib/mem.rb in mem-0.1.5

@@ -25,25 +25,41 @@
     memoized_table.delete(key)
   end
 
   module ClassMethods
     def memoize(method_name)
+      original_visibility =
+        case
+        when protected_instance_methods.include?(method_name)
+          :protected
+        when private_instance_methods.include?(method_name)
+          :private
+        else
+          :public
+        end
+
       define_method("#{method_name}_with_memoize") do |*args, &block|
         if has_memoized?(method_name)
           memoized(method_name)
         else
           memoize(method_name, send("#{method_name}_without_memoize", *args, &block))
         end
       end
+      send(original_visibility, "#{method_name}_with_memoize")
+
       alias_method "#{method_name}_without_memoize", method_name
       alias_method method_name, "#{method_name}_with_memoize"
 
       define_method("unmemoize_#{method_name}") do
         unmemoize(method_name)
       end
+      send(original_visibility, "unmemoize_#{method_name}")
 
-      define_method("#{method_name}=") do |value|
-        memoize(method_name, value)
+      if original_visibility != :private
+        define_method("#{method_name}=") do |value|
+          memoize(method_name, value)
+        end
+        send(original_visibility, "#{method_name}=")
       end
     end
   end
 end