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.

§