Friday, December 21, 2012

Python: Encapsulating Exceptions with Context Managers

Try/except can often interrupt the flow of logic making code harder to read. Take for example the following piece of code:

import sys

class Car(object):
   def create(self, color, stereo):
      try:
         vin = self._factory.make_car(color)
      except FactoryColorError, err:
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message), None, stacktrace

      try:
         self._customizer.update_car(vin, stereo)
      except CustomizerError, err:
         self._factory.remove_car(vin)
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message), None, stacktrace

      return vin

   def update(self, vin, color, stereo):
      try:
         self._factory.update_car_color(vin, color)
      except FactoryColorError, err:
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message), None, stacktrace
 
      try:
         self._customizer.update_car(vin, stereo)
      except CustomizerError, err:
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message), None, stacktrace

      return

A somewhat typical piece of code where the create and update are composed of multiple calls on the backend.  So the logic is as follows:

  • exception logic that changes third party exceptions into ValidationError
  • rollback exception logic in create that removes the car if the stereo option is invalid
The code is cluttered and the exception handling obscures the flow of both create and update.  The more exception logic we have the less readable the code is.  We also have some duplication of code. Enter context managers.

from contextlib import contextmanager
import sys

class Car(object):
   @contextmanager
   def _factory_error_handler(self):
      try:
         yield
      except FactoryColorError, err:
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message, None, stacktrace

   @contextmanager
   def _create_customizer_error_handler(self, vin):
      try:
         yield
      except CustomizerError, err:
         self._factory.remove_car(vin)
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message), None, stacktrace

   @contextmanager
   def _update_customizer_error_handler(self, vin):
      try:
         yield
      except CustomizerError, err:
         stacktrace = sys.exc_info()[2]
         raise ValidationError(err.message), None, stacktrace

   def create(self, color, stereo):
      with self._factory_error_handler():
         vin = self._factory.make_car(color)

      with self._create_customizer_error_handler(vin):
         self._customizer.update_car(vin, stereo)

      return vin

   def update(self, vin, color, stereo):
      with self._factory_error_handler():
         self._factory.update_car_color(vin, color)
 
      with self._update_customizer_error_handler(vin):
         self._customizer.update_car(vin, stereo)

      return


Refactoring the exception code to encapsulate the logic in it's own namespace in order to make the code more readable.  The core logic is no longer obscured by the exception handling, in fact the handlers make both the core and the exception logic easy to read and understand. We are also able to de-dupe some of the exception logic through reuse of the _factory_error_handler.  However, we can do better.

from contextlib import contextmanager
import sys

@contextmanager
def cleanup_error_handler(cleanup):
   try:
      yield
   except Exception:
      cleanup()
      raise

def make_error_handler(catch, throw):
   @contextmanager
   def handler():
      try:
         yield
      except catch, e:
         stacktrace = sys.exc_info()[2]
         raise throw(e.message), None, stacktrace
   return handler

factory_error_handler = make_error_handler(FactoryError, ValidationError)
customizer_error_handler = make_error_handler(CustomizerError, ValidationError)

class Car(object):
   def create(self, color, stereo):
      with factory_error_handler():
         vin = self._factory.make_car(color)

      cleanup = partial(self._factory.remove_car, vin)
      with cleanup_error_handler(cleanup), \
           customizer_error_handler():
         self._customizer.update_car(vin, stereo)

      return vin

   def update(self, vin, color, stereo):
      with factory_error_handler():
         self._factory.update_car_color(vin, color)
 
      with customizer_error_handler():
         self._customizer.update_car(vin, stereo)

      return


We retained the readability of using the handlers, but extracted them completely from the class.  The exception logic is now both encapsulated and decoupled from the specific class allowing reuse throughout the rest of the codebase.

§

Thursday, December 20, 2012

Python Metaprogramming: Dynamically Adding Methods to Classes

Dynamic addition of methods and attributes to classes and instances.

Dynamic Class Method Addition

Occasionally you may need dynamically add a function as a method to a class.  This is easily accomplished by assigning the function as an attribute of the class.


def fn(self):
  return id(self), self, type(self)

# Traditional Class Definition
class A_Class(object):
  def method_a(self):
    return id(self), self, type(self)

instance = A_Class()

# Modify the class and add fn as a method
setattr(A_Class, 'method_b', fn)

# Call the traditionally defined method
instance.method_a()
# Call the dynamically added method
instance.method_b()



Dynamic Instance Method Addition

When you add the method to a class all instances can access it.  If you only want a particular instance to have a method do this:

from types import MethodType

instance2 = A_Class()
setattr(instance, fn.__name__, MethodType(fn, instance, type(instance)))

# Calls the fn method attached to the instance
instance.fn()

# Throws an exception
instance2.fn()

§