Sunday, May 18, 2008

Metaprogramming Fun in Ruby

My friend, who is just learning Ruby, asked how he can automatically wrap some code around each user defined method in a class. The requirement (from his bosses) is to log when each method is being called.

To be more concrete, we want to factor out lines 3, 5, 9 and 11 from the following snippet.

    1 class AutoLogger
    2   def method_a
    3     puts "start logging method_a"
    4     puts "calling method_a"
    5     puts "end logging method_a"
    6   end
    7   
    8   def method_b
    9     puts "start logging method_b"
   10     puts "calling method_b"
   11     puts "end logging method_b"
   12   end
   13 end

The problem really boils down to two problems. First, how to redefine a method to be the original method body wrapped between the logging code. Second, how do we specify which methods in a class should be wrapped.

Each solution (except for the first one) relies on using Module#define_method to redefine the method in question. Sometimes we pass in a block that defines the original method, other times we make an UnboundMethod out of the original method and call it from the block that makes the new method body, binding it to self which at execution time is an instance of our class.

My friend said ideally, he wants every "user defined" method to automatically be wrapped.

The first solution I came up with was using method_missing. This solution has the drawback of having to refactor all your code to call log_some_method instead of some_method.

    1 class AutoLogger
    2   
    3   def method_missing(name, *args)
    4     name = name.to_s
    5     if name[0,4] == 'log_'
    6       method_name = name[4..-1]
    7       puts "start logging #{method_name}"
    8       self.send(method_name.to_sym, *args)
    9       puts "end logging #{method_name}"
   10     else
   11       super
   12     end
   13   end
   14   
   15   def method_a
   16     puts "calling method_a"
   17   end
   18   
   19   def method_b(arg1, arg2)
   20     puts "calling method_b(#{arg1}, #{arg2})"
   21   end
   22   
   23 end
   24 
   25 o = AutoLogger.new
   26 o.method_a
   27 o.method_b("test", 1)
   28 o.log_method_a
   29 o.log_method_b("test", 1)
   30 
   31 ##########
   32 # OUTPUT #
   33 ##########
   34 # calling method_a
   35 # calling method_b(test, 1)
   36 # start logging method_a
   37 # calling method_a
   38 # end logging method_a
   39 # start logging method_b
   40 # calling method_b(test, 1)
   41 # end logging method_b

With this approach, you can see that the original methods are kept intact and are callable. You have to call the special prefixed version of them to get the logging.

The next solution uses a special method to define methods instead of the def keyword.

    1 class AutoLogger
    2   
    3   def self.def_with_logging(method_name, &block)
    4     define_method(method_name) do |*args|
    5       puts "start logging #{method_name}"
    6       yield *args
    7       puts "end logging #{method_name}"
    8     end
    9   end
   10   
   11   def_with_logging(:method_a) do
   12     puts "calling method_a"
   13   end
   14   
   15   def_with_logging(:method_b) do |arg1, arg2|
   16     puts "calling method_b(#{arg1}, #{arg2})"
   17   end
   18   
   19 end
   20 
   21 o = AutoLogger.new
   22 o.method_a
   23 o.method_b("test", 1)
   24 
   25 ##########
   26 # OUTPUT #
   27 ##########
   28 # start logging method_a
   29 # calling method_a
   30 # end logging method_a
   31 # start logging method_b
   32 # calling method_b(test, 1)
   33 # end logging method_b

Again, this solution requires the programmer to refactor a lot of existing code (rewriting all your methods using def_with_logging instead of def).

The third solution is more "Rails-like" by adding a method to the class itself which can be called to add logging to existing methods.

    1 class AutoLogger
    2   @@inited = false
    3   
    4   def self.do_logging_for(*method_names)
    5     @@method_names ||= []
    6     @@method_names += method_names
    7   end
    8   
    9   do_logging_for :method_a, :method_b
   10   
   11   def initialize
   12     unless @@inited
   13       @@method_names.each do |method_name|
   14         self.class.class_eval do
   15           method = instance_method(method_name)
   16           define_method(method_name) do |*args|
   17             puts "start logging #{method_name}"
   18             method.bind(self).call(*args)
   19             puts "end logging #{method_name}"
   20           end
   21         end
   22       end
   23       @@inited = true
   24     end
   25   end
   26   
   27   def method_a
   28     puts "calling method_a"
   29   end
   30   
   31   def method_b(arg1, arg2)
   32     puts "calling method_b(#{arg1}, #{arg2})"
   33   end
   34   
   35 end
   36 
   37 o = AutoLogger.new
   38 o.method_a
   39 o.method_b("test", 1)
   40 
   41 ##########
   42 # OUTPUT #
   43 ##########
   44 # start logging method_a
   45 # calling method_a
   46 # end logging method_a
   47 # start logging method_b
   48 # calling method_b(test, 1)
   49 # end logging method_b

Notice how most of the work is deferred to initialize. This is so we can call do_logging_for before we define the methods it works on.

The fourth approach takes the least work because it does not require the programmer to do anything special or extra to get the desired effect. It all happens automatically.

    1 class AutoLogger
    2   
    3   def self.method_added(method_name)
    4     
    5     # so we don't get stuck in infinite recursion
    6     @@seen_methods ||= {}
    7     return if @@seen_methods.has_key?(method_name)
    8     @@seen_methods[method_name] = true
    9     
   10     method = instance_method(method_name)
   11     define_method(method_name) do |*args|
   12       puts "start logging #{method_name}"
   13       method.bind(self).call(*args)
   14       puts "end logging #{method_name}"
   15     end
   16   end
   17   
   18   def method_a
   19     puts "calling method_a"
   20   end
   21   
   22   def method_b(arg1, arg2)
   23     puts "calling method_b(#{arg1}, #{arg2})"
   24   end
   25   
   26 end
   27 
   28 o = AutoLogger.new
   29 o.method_a
   30 o.method_b("test", 1)
   31 
   32 ##########
   33 # OUTPUT #
   34 ##########
   35 # start logging method_a
   36 # calling method_a
   37 # end logging method_a
   38 # start logging method_b
   39 # calling method_b(test, 1)
   40 # end logging method_b

Class#method_added is a callback that gets fired, well, when you think it does. The interesting thing to note here is that we can go into infinite recursion because of how we redefinite each method. When a method is a defined the callback is triggered, in the callback we redefine the method which causes the callback to be triggered again.

That's it. Pretty neat what you can do with Ruby, huh?

1 comments:

jw said...

rubygasm