Writing context managers

Writing Functions in Python

Shayne Miel

Software Architect @ Duo Security

Two ways to define a context manager

  • Class-based
  • Function-based
Writing Functions in Python

Two ways to define a context manager

  • Class-based
  • Function-based *
Writing Functions in Python

How to create a context manager


def my_context():

# Add any set up code you need
yield
# Add any teardown code you need
  1. Define a function.
  2. (optional) Add any set up code your context needs.
  3. Use the "yield" keyword.
  4. (optional) Add any teardown code your context needs.
Writing Functions in Python

How to create a context manager

@contextlib.contextmanager
def my_context():

# Add any set up code you need
yield
# Add any teardown code you need
  1. Define a function.
  2. (optional) Add any set up code your context needs.
  3. Use the "yield" keyword.
  4. (optional) Add any teardown code your context needs.
  5. Add the `@contextlib.contextmanager` decorator.
Writing Functions in Python

The "yield" keyword

@contextlib.contextmanager
def my_context():
  print('hello')
  yield 42
  print('goodbye')
with my_context() as foo:
  print('foo is {}'.format(foo))
hello
foo is 42
goodbye
Writing Functions in Python

Setup and teardown

@contextlib.contextmanager
def database(url):
  # set up database connection
  db = postgres.connect(url)

  yield db

  # tear down database connection
  db.disconnect()
url = 'http://datacamp.com/data'
with database(url) as my_db:
  course_list = my_db.execute(
    'SELECT * FROM courses'
  )

blankspace

arrow

Writing Functions in Python

Setup and teardown

@contextlib.contextmanager
def database(url):
  # set up database connection
  db = postgres.connect(url)

  yield db

  # tear down database connection
  db.disconnect()
url = 'http://datacamp.com/data'
with database(url) as my_db:
  course_list = my_db.execute(
    'SELECT * FROM courses'
  )

blankspace

blankspace

blankspace

arrow

Writing Functions in Python

Setup and teardown

@contextlib.contextmanager
def database(url):
  # set up database connection
  db = postgres.connect(url)

  yield db

  # tear down database connection
  db.disconnect()
url = 'http://datacamp.com/data'
with database(url) as my_db:
  course_list = my_db.execute(
    'SELECT * FROM courses'
  )

blankspace

blankspace

blankspace

blankspace

blankspace

blankspace

arrow

Writing Functions in Python

Yielding a value or None

@contextlib.contextmanager
def database(url):
  # set up database connection
  db = postgres.connect(url)

  yield db

  # tear down database connection
  db.disconnect()
url = 'http://datacamp.com/data'
with database(url) as my_db:
  course_list = my_db.execute(
    'SELECT * FROM courses'
  )
@contextlib.contextmanager
def in_dir(path):
  # save current working directory
  old_dir = os.getcwd()

  # switch to new working directory
  os.chdir(path)

  yield

  # change back to previous
  # working directory
  os.chdir(old_dir)
with in_dir('/data/project_1/'):
  project_files = os.listdir()
Writing Functions in Python

Let's practice!

Writing Functions in Python

Preparing Video For Download...