DS in Python 4 — Context Managers

When writing Python programs, we often work with resources such as:

• Files
• Database connections
• Network connections
• Locks
• Temporary resources
• Memory-related resources

These resources usually need to be opened/acquired before use and closed/released after use.

For example, when working with a file:

file = open("data.txt")
data = file.read()
file.close()

We have to remember to close the file.

But what if an error occurs while reading the file?

file = open("data.txt")
data = file.read()  # Something goes wrong here
file.close()        # This may never execute

Now the file may not be closed properly.

Python provides context managers to solve this type of problem.

What is a Context Manager?

A context manager is a Python mechanism that manages a resource for you.

It makes sure that:

1. The resource is properly acquired before you use it.
2. The resource is properly released after you finish using it.
3. Cleanup happens even if an error occurs inside the block.

The most common way to use a context manager is with the with statement.

For example:

with open("data.txt") as file:
    data = file.read()

Here, Python automatically handles closing the file.

You don’t need to write:

file.close()

Why Do We Need Context Managers?

Imagine you are working with a file.

Without a context manager:

file = open("data.txt")
data = file.read()
file.close()

You are responsible for cleanup.

But consider:

file = open("data.txt")
data = file.read()
# Error occurs here

The program may exit the block before reaching:

file.close()

This can leave the resource open.

With a context manager:

with open("data.txt") as file:
    data = file.read()

Python takes care of cleanup.

Even if an exception occurs inside the with block, the context manager gets an opportunity to perform the required cleanup.

Components of a Context Manager

A class-based context manager generally has two important methods:

__enter__()
__exit__()

Let’s understand them.

1. __enter__()

__enter__() runs when Python enters the with block.

For example:

class MyContext:
    def __enter__(self):
        print("Entering context")
        return self
    def __exit__(self, exc_type, exc_value, traceback):
        print("Leaving context")

Now:

with MyContext():
    print("Inside context")

Output:

Entering context
Inside context
Leaving context

2. __exit__()

__exit__() runs when Python leaves the with block.

It runs even when an exception occurs inside the block, which is why it is useful for cleanup.

It receives three exception-related arguments:

__exit__(self, exc_type, exc_value, traceback)

They provide information about an exception, if one occurred.

Conceptually:

No exception:
exc_type  = None
exc_value = None
traceback = None

If an exception occurs, these contain information about that exception.

A Complete Custom Context Manager

Let’s create a simple example.

class MyContext:
    def __enter__(self):
        print("Resource acquired")
        return self
    def __exit__(self, exc_type, exc_value, traceback):
        print("Resource released")

with MyContext():
    print("Doing some work")

Output:

Resource acquired
Doing some work
Resource released

What Problem Does This Solve?

The main problem is:

How do we guarantee that a resource is cleaned up after we’re finished using it?

Examples:

Open file
   ↓
Read file
   ↓
Close file

Open database connection
   ↓
Execute query
   ↓
Close connection

Acquire lock
   ↓
Do work
   ↓
Release lock

Context managers provide a clean and reliable way to handle these situations.

The with Statement

The with statement is the most common way to use a context manager.

For example:

with open("data.csv") as file:
    data = file.read()

Conceptually, Python is doing something similar to:

Enter context
      ↓
Open/acquire resource
      ↓
Run your code
      ↓
Exit context
      ↓
Clean up resource

This is safer than manually managing the resource.

Context Managers in Data Science

Context managers are particularly useful in Data Science because Data Science applications often work with external resources.

Common examples include:

Files
Databases
Temporary files
Cloud resources
Locks
Experimental resources

Let’s look at some practical examples.

1. Reading Large Data Files

Suppose you’re working with a large CSV:

with open("large_dataset.csv") as file:
    for line in file:
        process(line)

The file is opened when entering the context.

When the block finishes, Python closes the file.

CSV file
   ↓
Open
   ↓
Process data
   ↓
Finish
   ↓
Automatically close

This is especially useful when processing large datasets line by line.

2. Reading Data with Pandas

You may also encounter context managers when working with Pandas-related file operations.

For example:

with open("data.csv") as file:
    # read/process file
    pass

The context manager ensures that the underlying file resource is cleaned up.

For many Pandas functions, you don’t manually manage the file because Pandas handles the resource internally. But understanding context managers helps you understand what’s happening when Python code interacts with files and other resources.

3. Database Connections

Data Scientists frequently query databases.

Conceptually:

with database_connection() as connection:
    data = connection.execute(
        "SELECT * FROM customers"
    )

The flow becomes:

Database
   ↓
Open connection
   ↓
Execute query
   ↓
Get data
   ↓
Close/release connection

The exact context-manager behavior depends on the database library you’re using, but the pattern is very common.

4. Temporary Files

Suppose you need a temporary file while processing data.

A context manager can make sure the temporary resource is cleaned up afterward.

For example, Python provides TemporaryFile:

from tempfile import TemporaryFile

with TemporaryFile() as file:
    file.write(b"some data")
    file.seek(0)
    data = file.read()

After the with block finishes, the temporary file is automatically closed and handled according to the temporary-file API.

Get Abhishek Jain’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

This can be useful when a Data Science workflow needs intermediate data without permanently storing it.

5. Managing Memory or Resources

Data Science applications can consume substantial resources.

For example:

Load large dataset
      ↓
Process dataset
      ↓
Release resource

Context managers can be used to make resource acquisition and cleanup explicit and reliable.

However, an important distinction is:

A context manager doesn’t magically free all Python memory.

Python’s garbage collector and reference counting manage Python objects. A context manager is specifically about entering and exiting a managed context, often to acquire and release an external resource.

6. Timing Data Science Operations

You can also create a custom context manager to measure execution time.

import time

class Timer:
    def __enter__(self):
        self.start = time.time()
    def __exit__(self, exc_type, exc_value, traceback):
        end = time.time()
        print("Time:", end - self.start)

Use it like this:

with Timer():
    # Data processing
    for i in range(1_000_000):
        pass

Output might be:

Time: 0.05

This can be useful when comparing:

Data cleaning
Feature engineering
Model training
Data loading

and identifying which part of a pipeline is taking the most time.

Context Managers with contextlib

Python also provides a module called contextlib that makes creating context managers easier.

One useful tool is:

from contextlib import contextmanager

We can create a context manager using a generator function.

For example:

from contextlib import contextmanager

@contextmanager
def my_context():
    print("Start")
    yield
    print("End")

with my_context():
    print("Doing work")

Output:

Start
Doing work
End

Here, yield separates the entering part from the exiting part.

This is similar in spirit to the __enter__() / __exit__() approach, but it can be much simpler for straightforward context managers.

Context Manager vs Decorator

These concepts can sometimes be confused because both can add behavior around code.

A decorator works mainly around a function call:

Decorator
    ↓
Function

A context manager works around a block of code:

Context Manager
      ↓
with block

For example:

@timer
def train_model():
    pass

versus:

with Timer():
    train_model()

Both could measure time, but they use different mechanisms.

Advantages of Context Managers

1. Automatic Cleanup

Resources are cleaned up when the context is exited.

For example:

with open("data.csv") as file:
    ...

You don’t have to remember:

file.close()

2. Safer Code

Cleanup happens even if an exception occurs inside the with block.

This makes resource management more reliable.

3. Cleaner Code

Compare:

file = open("data.csv")

try:
    data = file.read()
finally:
    file.close()

with:

with open("data.csv") as file:
    data = file.read()

The second version is easier to read.

4. Separation of Responsibilities

Your main code focuses on the actual work.

The context manager handles:

Acquire
Use
Release

This makes large applications easier to organize.

5. Reusable Resource Management

Once you create a context manager, you can use it in many places.

For example:

with Timer():
    clean_data()
with Timer():
    train_model()
with Timer():
    evaluate_model()

The timing logic doesn’t need to be repeated.

Disadvantages of Context Managers

1. Can Be Confusing Initially

Concepts such as:

__enter__()
__exit__()

and:

@contextmanager
yield

can be confusing when you’re first learning Python.

2. Not Needed Everywhere

You don’t need a context manager for every piece of code.

For simple calculations:

x = 10
y = 20
result = x + y

there is no resource that needs to be acquired and released.

Using a context manager would add unnecessary complexity.

3. Custom Context Managers Require Care

If you create your own context manager, you need to make sure cleanup behavior is correct.

A poorly designed context manager can hide errors or fail to release resources properly.

A Real Data Science Example

Imagine you are building a data-processing pipeline:

Load data
    ↓
Clean data
    ↓
Feature engineering
    ↓
Train model
    ↓
Evaluate model

You want to measure the time spent in each stage.

You could create:

from contextlib import contextmanager
import time

@contextmanager
def timer(name):
    start = time.time()
    yield
    end = time.time()
    print(f"{name}: {end - start:.2f} seconds")

with timer("Data cleaning"):
    clean_data()
with timer("Feature engineering"):
    create_features()
with timer("Model training"):
    train_model()

Output might look like:

Data cleaning: 2.13 seconds
Feature engineering: 5.42 seconds
Model training: 18.67 seconds

This gives you a simple way to monitor different stages of a Data Science pipeline.

Context Managers and Resource Management

The biggest takeaway is that context managers are fundamentally about managing a context and guaranteeing exit/cleanup behavior.

A useful mental model is:

        ENTER
          ↓
   Acquire / Prepare
          ↓
     DO THE WORK
          ↓
        EXIT
          ↓
   Release / Cleanup

For example:

File:
Open → Read → Close

Database:
Connect → Query → Close

Lock:
Acquire → Work → Release

Temporary resource:
Create → Use → Clean up

Final Takeaway

A context manager provides a structured way to manage something before, during, and after a block of code.

The most common syntax is:

with something:
    # do work

The two important methods for a class-based context manager are:

__enter__()
__exit__()

And Python’s contextlib module provides tools such as @contextmanager for creating context managers more conveniently.

For Data Science, context managers are useful when working with:

• Large files
• Database connections
• Temporary resources
• Timing/profiling
• Resource management
• Data-processing pipelines

The easiest way to remember the concept is:

A context manager says: “I’ll prepare this resource before you use it, and I’ll make sure the necessary cleanup happens when you’re done.”