Snippets

Lilian Besson (Naereen) Python context manager: conditionally executing body?

Created by Lilian Besson (Naereen) last modified Lilian Besson
*.pyc
__pycache__
class ProbabilityError(Exception):
     pass 

class probability(object):
     """ Context manager to be called like this:

     >>> with probability(0.5):
     ...     print("This happens once every two trials.")
     """

     def __init__(self, probability=0.5):
          assert 0 <= probability <= 1, "Error: probability parameter should be in [0, 1]."  # DEBUG
          self.probability = probability
          from random import random
          self.random = random

     def __enter__(self):
          if self.random() >= self.probability:
               print("Exiting...")  # DEBUG
               raise ProbabilityError  # <--- this fails!
          else:
               print("Doing...")  # DEBUG
               return self

     def __exit__(self, *exc):
          print("*exc =", exc)
          if exc and exc[0] == ProbabilityError:
               print("I saw ProbabilityError...")  # DEBUG
               return True

Comments (0)