Python Decorators and Context Managers: Closures, @wraps, Parameterized Decorators, Class Decorators, the with Statement, Custom Context Managers, contextlib, and Every Pattern Data Engineers Need
You have written retry logic for the third time this week. You have added timing code to 12 functions. You have wrapped 8 database calls in try/finally blocks to close connections. Every function has the same boilerplate — the same try/except, the same logging, the same cleanup. The actual business logic is buried under layers of repetitive infrastructure code.
Decorators and context managers exist to eliminate this repetition. A decorator wraps a function with reusable behavior — retry, timing, logging, caching, validation — without changing the function itself. A context manager guarantees setup and cleanup — open a connection, do work, close the connection — even if an error occurs. Together, they are the two most powerful patterns for writing clean, production-grade Python.
Real-life analogy: A decorator is like a security checkpoint at an airport. You (the function) walk through the checkpoint (the decorator), which adds scanning, ID verification, and boarding pass checks — without changing who you are or where you are going. Every passenger goes through the same checkpoint, but each passenger’s journey (function logic) is different. The checkpoint is reusable across all passengers. A context manager is like a hotel room — you check in (setup), use the room (do your work), and check out (cleanup). Whether you leave normally or the fire alarm goes off, the room is cleaned and made available for the next guest. The with statement is your reservation system.
Table of Contents
- Part 1: Decorators
- Functions Are Objects
- Closures — Functions That Remember
- Your First Decorator
- The @syntax Sugar
- @functools.wraps — Preserving Function Identity
- Decorators with Arguments
- Common Built-in Decorators
- @staticmethod, @classmethod, @property (Recap)
- @functools.lru_cache — Memoization
- @functools.total_ordering
- Stacking Multiple Decorators
- Class-Based Decorators
- Data Engineering Decorator Patterns
- @timer — Measure Execution Time
- @retry — Retry with Exponential Backoff
- @log_call — Log Function Calls
- @validate_input — Input Validation
- @deprecated — Mark Functions as Deprecated
- Part 2: Context Managers
- The with Statement Recap
- How Context Managers Work (__enter__ and __exit__)
- Custom Context Managers with Classes
- contextlib.contextmanager — The Easy Way
- Nested Context Managers
- contextlib.suppress — Ignore Specific Exceptions
- Data Engineering Context Manager Patterns
- Database Connection Manager
- Pipeline Timer
- Temporary Directory Manager
- File Lock Manager
- Audit Logger Context Manager
- Combining Decorators and Context Managers
- Common Mistakes
- Interview Questions
- Wrapping Up
Part 1: Decorators
Functions Are Objects
Before understanding decorators, you need to understand that functions in Python are first-class objects. They can be assigned to variables, passed as arguments, returned from other functions, and stored in data structures — just like integers or strings.
# Functions can be assigned to variables
def greet(name):
return f"Hello, {name}!"
say_hello = greet # Assign function to a variable (no parentheses!)
print(say_hello("Naveen")) # "Hello, Naveen!" — same function, different name
# Functions can be passed as arguments
def apply(func, value):
return func(value)
print(apply(greet, "Alice")) # "Hello, Alice!"
print(apply(len, [1, 2, 3])) # 3
print(apply(str.upper, "hello")) # "HELLO"
# Functions can be returned from other functions
def create_multiplier(n):
def multiply(x):
return x * n
return multiply # Return the inner function
double = create_multiplier(2)
triple = create_multiplier(3)
print(double(5)) # 10
print(triple(5)) # 15
# Functions have attributes
print(greet.__name__) # "greet"
print(greet.__doc__) # None (no docstring in this example)Closures — Functions That Remember
A closure is a function that remembers the variables from its enclosing scope, even after that scope has finished executing. This is the foundation of decorators.
def make_counter():
count = 0 # Variable in the enclosing scope
def increment():
nonlocal count # "I want to modify count from the outer scope"
count += 1
return count
return increment # Return the inner function
# The inner function "remembers" count even after make_counter() has returned
counter = make_counter()
print(counter()) # 1
print(counter()) # 2
print(counter()) # 3
# Each call to make_counter() creates a SEPARATE closure with its OWN count
counter_a = make_counter()
counter_b = make_counter()
print(counter_a()) # 1
print(counter_a()) # 2
print(counter_b()) # 1 — counter_b has its own count, starting from 0Real-life analogy: A closure is like a factory that produces robots. Each robot (increment) carries its own battery (count). When the factory (make_counter) shuts down, the robots keep running with their batteries. Each robot tracks its own battery level independently.
Your First Decorator
# A decorator is a function that takes a function and returns a NEW function
# The new function usually calls the original and adds behavior before/after
import time
def timer(func):
"""Decorator that measures how long a function takes."""
def wrapper(*args, **kwargs):
start = time.time()
result = func(*args, **kwargs) # Call the ORIGINAL function
elapsed = time.time() - start
print(f"{func.__name__} took {elapsed:.2f}s")
return result # Return the original result
return wrapper
# Apply the decorator MANUALLY
def load_customers():
time.sleep(1) # Simulate work
return [{"id": 1, "name": "Naveen"}]
load_customers = timer(load_customers) # Replace function with wrapped version
result = load_customers()
# "load_customers took 1.00s"
# What happened:
# 1. timer(load_customers) created wrapper() — which remembers load_customers via closure
# 2. We replaced load_customers with wrapper
# 3. Calling load_customers() now calls wrapper(), which:
# a. Records start time
# b. Calls the ORIGINAL load_customers()
# c. Records end time
# d. Prints the duration
# e. Returns the original resultThe @syntax Sugar
# The @ syntax is just shorthand for the manual pattern above
@timer # Equivalent to: load_customers = timer(load_customers)
def load_customers():
time.sleep(1)
return [{"id": 1, "name": "Naveen"}]
result = load_customers() # Automatically timed
# "load_customers took 1.00s"
# @ is syntactic sugar — nothing more. These are identical:
# @decorator == func = decorator(func)
# def func(): ...
# The decorator is applied ONCE at function definition time, not at call time@functools.wraps — Preserving Function Identity
import functools
import time
# WITHOUT @wraps — the decorated function loses its identity
def timer_bad(func):
def wrapper(*args, **kwargs):
start = time.time()
result = func(*args, **kwargs)
print(f"{func.__name__} took {time.time() - start:.2f}s")
return result
return wrapper
@timer_bad
def load_customers():
"""Load customer data from the database."""
pass
print(load_customers.__name__) # "wrapper" — WRONG! Should be "load_customers"
print(load_customers.__doc__) # None — WRONG! Lost the docstring
# WITH @wraps — the decorated function keeps its identity
def timer_good(func):
@functools.wraps(func) # Copy __name__, __doc__, etc. from func to wrapper
def wrapper(*args, **kwargs):
start = time.time()
result = func(*args, **kwargs)
print(f"{func.__name__} took {time.time() - start:.2f}s")
return result
return wrapper
@timer_good
def load_customers():
"""Load customer data from the database."""
pass
print(load_customers.__name__) # "load_customers" — CORRECT!
print(load_customers.__doc__) # "Load customer data from the database." — CORRECT!
# RULE: ALWAYS use @functools.wraps(func) in every decorator you write.
# Without it, debugging tools, documentation generators, and serialization
# frameworks cannot identify your functions.Decorators with Arguments
Sometimes you need to configure a decorator — “retry 5 times” vs “retry 3 times.” This requires a decorator factory — a function that returns a decorator.
import functools
import time
# A decorator factory — takes arguments and RETURNS a decorator
def retry(max_attempts=3, delay=1, exceptions=(Exception,)):
"""Decorator that retries a function on failure."""
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
for attempt in range(1, max_attempts + 1):
try:
return func(*args, **kwargs)
except exceptions as e:
if attempt == max_attempts:
raise
wait = delay * (2 ** (attempt - 1))
print(f" Attempt {attempt} failed: {e}. Retrying in {wait}s...")
time.sleep(wait)
return wrapper
return decorator
# Usage WITH arguments
@retry(max_attempts=5, delay=2, exceptions=(ConnectionError, TimeoutError))
def fetch_data(url):
"""Fetch data from an API."""
response = requests.get(url, timeout=10)
return response.json()
# Usage WITH defaults (still need parentheses!)
@retry()
def load_table(name):
return query(f"SELECT * FROM {name}")
# The three layers:
# retry(max_attempts=5) → returns decorator
# decorator(fetch_data) → returns wrapper
# wrapper(*args) → calls fetch_data with retry logicReal-life analogy: A decorator factory is like a vending machine that produces security checkpoints. You configure the machine — “check IDs, allow 3 retries, scan for metal.” It produces a checkpoint with those specific settings. Each checkpoint (decorator) is then installed at a gate (applied to a function). Different gates can have different checkpoint configurations from the same machine.
Common Built-in Decorators
@staticmethod, @classmethod, @property (Recap)
# Quick recap from the OOP post — these are all decorators!
class Pipeline:
@staticmethod # No self or cls — utility function
def is_valid_table(name):
return bool(name and not name.startswith("_"))
@classmethod # Receives cls, not self — alternative constructor
def from_config(cls, config_path):
config = load_json(config_path)
return cls(config["name"], config["source"])
@property # Access method as an attribute
def status(self):
return "running" if self._active else "stopped"@functools.lru_cache — Memoization
import functools
# lru_cache caches function results — same input = cached output (no recomputation)
# LRU = Least Recently Used — evicts oldest entries when cache is full
@functools.lru_cache(maxsize=128) # Cache up to 128 unique calls
def get_table_schema(table_name):
"""Expensive operation — query the database for column metadata."""
print(f" Querying schema for {table_name}...") # Only prints on cache MISS
return query(f"SELECT column_name, data_type FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = '{table_name}'")
# First call — cache MISS, queries the database
schema = get_table_schema("customers")
# Querying schema for customers...
# Second call — cache HIT, returns instantly (no database query)
schema = get_table_schema("customers")
# (nothing printed — result served from cache)
# Different argument — cache MISS
schema = get_table_schema("orders")
# Querying schema for orders...
# Check cache statistics
print(get_table_schema.cache_info())
# CacheInfo(hits=1, misses=2, maxsize=128, currsize=2)
# Clear the cache
get_table_schema.cache_clear()
# IMPORTANT: Arguments must be HASHABLE (no lists or dicts)
# Use @functools.cache for unlimited cache (Python 3.9+)Real-life analogy: lru_cache is like a speed dial on a phone. The first time you call a number, you dial all 10 digits (compute the result). Then you save it to speed dial (cache). Next time, you press one button and connect instantly (cache hit). If your speed dial is full (maxsize), the oldest unused number is removed to make room.
@functools.total_ordering
import functools
# total_ordering generates ALL comparison methods from just __eq__ and ONE of __lt__/__gt__
@functools.total_ordering
class PipelineResult:
def __init__(self, table_name, row_count):
self.table_name = table_name
self.row_count = row_count
def __eq__(self, other):
return self.row_count == other.row_count
def __lt__(self, other):
return self.row_count < other.row_count
# Now ALL comparisons work: ==, !=, <, >, <=, >=
r1 = PipelineResult("customers", 50000)
r2 = PipelineResult("orders", 100000)
print(r1 < r2) # True
print(r1 > r2) # False (auto-generated!)
print(r1 <= r2) # True (auto-generated!)
print(r1 >= r2) # False (auto-generated!)
print(sorted([r2, r1])) # Sorted by row_countStacking Multiple Decorators
# Decorators are applied BOTTOM to TOP (closest to the function first)
@log_call # Applied THIRD (outermost)
@timer # Applied SECOND
@retry(max_attempts=3) # Applied FIRST (closest to function)
def load_customers():
return fetch_from_database("customers")
# Equivalent to:
# load_customers = log_call(timer(retry(max_attempts=3)(load_customers)))
# Execution order (when called):
# 1. log_call's wrapper runs first (logs the call)
# 2. timer's wrapper runs (starts timing)
# 3. retry's wrapper runs (handles retries)
# 4. load_customers() runs (the actual function)
# 3. retry returns (or retries on failure)
# 2. timer records elapsed time
# 1. log_call finishes loggingReal-life analogy: Stacking decorators is like putting on layers of clothing. The function is you. The innermost decorator (@retry) is your shirt — closest to you. The next (@timer) is your jacket. The outermost (@log_call) is your coat. When you get dressed, you put on the shirt first (bottom decorator applied first). When someone sees you, they see the coat first (outermost decorator executes first).
Class-Based Decorators
# Use a class as a decorator when you need to maintain state across calls
class CallCounter:
"""Decorator that counts how many times a function is called."""
def __init__(self, func):
functools.update_wrapper(self, func)
self.func = func
self.count = 0
def __call__(self, *args, **kwargs):
self.count += 1
print(f"{self.func.__name__} called {self.count} time(s)")
return self.func(*args, **kwargs)
@CallCounter
def process_table(name):
print(f"Processing {name}")
process_table("customers") # "process_table called 1 time(s)"
process_table("orders") # "process_table called 2 time(s)"
print(process_table.count) # 2 — access the stateData Engineering Decorator Patterns
@timer — Measure Execution Time
import functools
import time
import logging
logger = logging.getLogger(__name__)
def timer(func):
"""Log the execution time of the decorated function."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
start = time.time()
try:
result = func(*args, **kwargs)
elapsed = time.time() - start
level = logging.WARNING if elapsed > 60 else logging.INFO
logger.log(level, f"{func.__name__} completed in {elapsed:.2f}s")
return result
except Exception as e:
elapsed = time.time() - start
logger.error(f"{func.__name__} FAILED after {elapsed:.2f}s: {e}")
raise
return wrapper
@timer
def extract_customers():
# ... extraction logic
time.sleep(2)
return [{"id": 1}]
extract_customers()
# INFO: extract_customers completed in 2.00s
# Auto-warns if slow (> 60 seconds):
# WARNING: extract_large_table completed in 125.30s@retry — Retry with Exponential Backoff
import functools
import time
import logging
logger = logging.getLogger(__name__)
def retry(max_attempts=3, delay=1, backoff=2, exceptions=(Exception,)):
"""Retry a function with configurable exponential backoff."""
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
for attempt in range(1, max_attempts + 1):
try:
return func(*args, **kwargs)
except exceptions as e:
if attempt == max_attempts:
logger.error(f"{func.__name__}: all {max_attempts} attempts failed")
raise
wait = delay * (backoff ** (attempt - 1))
logger.warning(
f"{func.__name__}: attempt {attempt}/{max_attempts} failed: {e}. "
f"Retrying in {wait}s..."
)
time.sleep(wait)
return wrapper
return decorator
@retry(max_attempts=4, delay=2, exceptions=(ConnectionError, TimeoutError))
def fetch_api_data(url):
"""Fetch data from an external API."""
response = requests.get(url, timeout=30)
response.raise_for_status()
return response.json()
# Attempt 1: fails → wait 2s
# Attempt 2: fails → wait 4s
# Attempt 3: fails → wait 8s
# Attempt 4: succeeds (or raises the exception)@log_call — Log Function Calls
import functools
import logging
logger = logging.getLogger(__name__)
def log_call(func):
"""Log every call to the decorated function with arguments and return value."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
args_repr = [repr(a) for a in args]
kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
signature = ", ".join(args_repr + kwargs_repr)
logger.info(f"Calling {func.__name__}({signature})")
try:
result = func(*args, **kwargs)
logger.info(f"{func.__name__} returned {result!r}")
return result
except Exception as e:
logger.exception(f"{func.__name__} raised {type(e).__name__}: {e}")
raise
return wrapper
@log_call
def load_table(table_name, batch_size=10000):
return {"rows": 50000, "status": "success"}
load_table("customers", batch_size=25000)
# INFO: Calling load_table('customers', batch_size=25000)
# INFO: load_table returned {'rows': 50000, 'status': 'success'}@validate_input — Input Validation
import functools
def validate_input(**validators):
"""Validate function arguments using provided validators.
Usage: @validate_input(name=str, age=lambda x: 0 < x < 150)
"""
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
# Get the function's parameter names
import inspect
sig = inspect.signature(func)
bound = sig.bind(*args, **kwargs)
bound.apply_defaults()
for param_name, validator in validators.items():
if param_name in bound.arguments:
value = bound.arguments[param_name]
if isinstance(validator, type):
if not isinstance(value, validator):
raise TypeError(
f"{param_name} must be {validator.__name__}, "
f"got {type(value).__name__}"
)
elif callable(validator):
if not validator(value):
raise ValueError(
f"Validation failed for {param_name}: {value!r}"
)
return func(*args, **kwargs)
return wrapper
return decorator
@validate_input(
table_name=str,
batch_size=lambda x: isinstance(x, int) and x > 0,
environment=lambda x: x in ("dev", "staging", "production")
)
def load_table(table_name, batch_size=10000, environment="production"):
print(f"Loading {table_name} in {environment} with batch {batch_size}")
load_table("customers", batch_size=50000) # Works
load_table("customers", batch_size=-1) # ValueError
load_table("customers", environment="invalid") # ValueError
load_table(123) # TypeError@deprecated — Mark Functions as Deprecated
import functools
import warnings
def deprecated(reason=""):
"""Mark a function as deprecated with an optional reason."""
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
message = f"{func.__name__} is deprecated."
if reason:
message += f" {reason}"
warnings.warn(message, DeprecationWarning, stacklevel=2)
return func(*args, **kwargs)
return wrapper
return decorator
@deprecated("Use load_table_v2() instead.")
def load_table(name):
return query(f"SELECT * FROM {name}")
load_table("customers")
# DeprecationWarning: load_table is deprecated. Use load_table_v2() instead.Part 2: Context Managers
The with Statement Recap
# The with statement guarantees cleanup — even if an error occurs
# File handling (most common context manager)
with open("data.csv", "r") as f:
data = f.read()
# f.close() is called automatically — even if f.read() raises an error
# Database connection
with get_connection() as conn:
conn.execute("INSERT INTO ...")
# conn.close() is called automatically
# The with statement calls TWO methods on the object:
# __enter__() — runs at the START (setup: open file, connect to DB)
# __exit__() — runs at the END (cleanup: close file, close connection)
# __exit__ runs even if an exception occurs inside the with blockHow Context Managers Work (__enter__ and __exit__)
# What happens when you write:
# with SomeContext() as x:
# do_something(x)
# Python translates this to:
# 1. context = SomeContext()
# 2. x = context.__enter__() ← setup
# 3. try:
# do_something(x)
# except Exception as e:
# if not context.__exit__(type(e), e, traceback):
# raise ← re-raise if __exit__ returns False
# else:
# context.__exit__(None, None, None) ← no exception
# __enter__ returns the object to bind to "as x"
# __exit__ receives exception info (or None if no exception)
# __exit__ returns True to SUPPRESS the exception, False to let it propagateCustom Context Managers with Classes
import time
import logging
logger = logging.getLogger(__name__)
class PipelineTimer:
"""Context manager that times a block of code and logs the result."""
def __init__(self, step_name, warn_threshold=60):
self.step_name = step_name
self.warn_threshold = warn_threshold
self.start = None
self.elapsed = None
def __enter__(self):
logger.info(f"[{self.step_name}] Started")
self.start = time.time()
return self # The "as" variable — you can access self.elapsed later
def __exit__(self, exc_type, exc_val, exc_tb):
self.elapsed = time.time() - self.start
if exc_type:
logger.error(f"[{self.step_name}] FAILED after {self.elapsed:.2f}s: {exc_val}")
elif self.elapsed > self.warn_threshold:
logger.warning(f"[{self.step_name}] SLOW: {self.elapsed:.2f}s (threshold: {self.warn_threshold}s)")
else:
logger.info(f"[{self.step_name}] Completed in {self.elapsed:.2f}s")
return False # Do NOT suppress exceptions — let them propagate
# Usage
with PipelineTimer("extract_customers") as timer:
data = extract_from_database("customers")
transform(data)
print(f"Total time: {timer.elapsed:.2f}s") # Access elapsed time after the block
# Output:
# INFO: [extract_customers] Started
# INFO: [extract_customers] Completed in 3.25scontextlib.contextmanager — The Easy Way
Writing a full class with __enter__ and __exit__ is verbose. The contextlib.contextmanager decorator lets you write context managers as generator functions — much simpler.
from contextlib import contextmanager
import time
import logging
logger = logging.getLogger(__name__)
@contextmanager
def pipeline_timer(step_name, warn_threshold=60):
"""Context manager using a generator — much simpler than a class."""
logger.info(f"[{step_name}] Started")
start = time.time()
try:
yield start # Everything BEFORE yield = __enter__
# The yield value becomes the "as" variable
except Exception as e:
elapsed = time.time() - start
logger.error(f"[{step_name}] FAILED after {elapsed:.2f}s: {e}")
raise # Re-raise the exception
else:
elapsed = time.time() - start
level = logging.WARNING if elapsed > warn_threshold else logging.INFO
logger.log(level, f"[{step_name}] Completed in {elapsed:.2f}s")
finally:
pass # Everything AFTER yield = __exit__
# Runs whether or not an exception occurred
# Usage — identical to the class version
with pipeline_timer("extract_customers"):
data = extract_from_database("customers")
# The pattern:
# 1. Code BEFORE yield = setup (__enter__)
# 2. yield = pause and give control to the with block
# 3. Code AFTER yield = cleanup (__exit__)Real-life analogy: The @contextmanager decorator turns a function into a hotel room manager. Everything before yield is the check-in process (prepare the room). yield is handing over the room key to the guest. Everything after yield is the checkout process (clean the room). Whether the guest leaves normally or the fire alarm goes off, checkout always happens.
Nested Context Managers
# Multiple context managers in one with statement
with open("input.csv", "r") as src, open("output.csv", "w") as dst:
dst.write(src.read())
# Equivalent (for Python 3.10+, multi-line with parentheses):
with (
open("input.csv", "r") as src,
open("output.csv", "w") as dst,
pipeline_timer("copy_file") as timer,
):
dst.write(src.read())
# Nested with statements (for complex setups)
with DatabaseConnection("source") as source_conn:
with DatabaseConnection("target") as target_conn:
data = source_conn.query("SELECT * FROM customers")
target_conn.insert("silver.customers", data)contextlib.suppress — Ignore Specific Exceptions
from contextlib import suppress
# Instead of try/except with pass:
try:
os.remove("temp.csv")
except FileNotFoundError:
pass
# Use suppress — cleaner:
with suppress(FileNotFoundError):
os.remove("temp.csv")
# Multiple exception types:
with suppress(FileNotFoundError, PermissionError):
os.remove("locked_file.csv")
# Practical: delete temp files without caring if they exist
temp_files = ["temp_1.csv", "temp_2.csv", "staging.parquet"]
for f in temp_files:
with suppress(FileNotFoundError):
os.remove(f)Data Engineering Context Manager Patterns
Database Connection Manager
from contextlib import contextmanager
import logging
logger = logging.getLogger(__name__)
@contextmanager
def database_connection(connection_string, autocommit=False):
"""Manage a database connection with automatic commit/rollback/close."""
conn = None
try:
conn = create_connection(connection_string)
conn.autocommit = autocommit
logger.info("Database connection opened")
yield conn
except Exception as e:
if conn and not autocommit:
conn.rollback()
logger.error(f"Transaction rolled back: {e}")
raise
else:
if conn and not autocommit:
conn.commit()
logger.info("Transaction committed")
finally:
if conn:
conn.close()
logger.info("Database connection closed")
# Usage — connection is ALWAYS closed, transactions are ALWAYS handled
with database_connection(CONN_STRING) as conn:
conn.execute("INSERT INTO audit_log VALUES (...)")
conn.execute("UPDATE customers SET status = 'active' WHERE ...")
# Auto-committed on success, auto-rolled-back on failure, auto-closed alwaysTemporary Directory Manager
from contextlib import contextmanager
from pathlib import Path
import tempfile
import shutil
@contextmanager
def temp_working_dir(prefix="pipeline_"):
"""Create a temporary directory that is automatically cleaned up."""
temp_dir = Path(tempfile.mkdtemp(prefix=prefix))
logger.info(f"Created temp directory: {temp_dir}")
try:
yield temp_dir
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
logger.info(f"Cleaned up temp directory: {temp_dir}")
# Usage — temp files are automatically cleaned up
with temp_working_dir("etl_") as work_dir:
# Download files to temp directory
input_file = work_dir / "raw_data.csv"
download(url, input_file)
# Process
output_file = work_dir / "clean_data.parquet"
transform(input_file, output_file)
# Upload to data lake
upload(output_file, "bronze/customers/")
# work_dir and all its contents are deleted here — no manual cleanupAudit Logger Context Manager
from contextlib import contextmanager
from datetime import datetime, timezone
import json
@contextmanager
def audit_pipeline(pipeline_name, config=None):
"""Track pipeline execution with structured audit logging."""
run_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
start = datetime.now(timezone.utc)
audit_record = {
"pipeline": pipeline_name,
"run_id": run_id,
"started_at": start.isoformat(),
"config": config or {},
"status": "running",
}
logger.info(f"Pipeline '{pipeline_name}' started (run_id: {run_id})")
try:
yield audit_record # The pipeline code can update this dict
except Exception as e:
audit_record["status"] = "failed"
audit_record["error"] = str(e)
logger.error(f"Pipeline '{pipeline_name}' FAILED: {e}")
raise
else:
audit_record["status"] = "success"
finally:
audit_record["completed_at"] = datetime.now(timezone.utc).isoformat()
audit_record["duration_seconds"] = (
datetime.now(timezone.utc) - start
).total_seconds()
# Write audit record to file or database
with open(f"audit/{run_id}.json", "w") as f:
json.dump(audit_record, f, indent=2)
logger.info(
f"Pipeline '{pipeline_name}' {audit_record['status']} "
f"in {audit_record['duration_seconds']:.1f}s"
)
# Usage
with audit_pipeline("daily_customer_load", config={"source": "azure_sql"}) as audit:
data = extract("customers")
audit["rows_extracted"] = len(data) # Add metrics to audit record
clean = transform(data)
audit["rows_transformed"] = len(clean)
load(clean, "silver.customers")
audit["rows_loaded"] = len(clean)
# Audit JSON is written automatically — success or failureCombining Decorators and Context Managers
# The ultimate production pattern: decorators for function-level concerns,
# context managers for resource-level concerns
@timer # Time the whole function
@retry(max_attempts=3, exceptions=(ConnectionError,)) # Retry on connection failure
@log_call # Log every call
def load_daily_customers():
"""Full daily customer load with automatic resource management."""
with database_connection(SOURCE_CONN) as source: # Auto-close source
with database_connection(TARGET_CONN) as target: # Auto-close target
with pipeline_timer("extract"): # Time extraction
data = source.query("SELECT * FROM customers")
with pipeline_timer("transform"): # Time transformation
clean = transform(data)
with pipeline_timer("load"): # Time load
target.insert("silver.customers", clean)
return {"rows": len(clean), "status": "success"}
# This function automatically:
# ✅ Logs every call with arguments (@ log_call)
# ✅ Retries up to 3 times on ConnectionError (@retry)
# ✅ Measures total execution time (@timer)
# ✅ Opens and closes database connections (with database_connection)
# ✅ Times each pipeline step (with pipeline_timer)
# ✅ Rolls back on failure, commits on success
# All without a single try/except or close() call in the business logicCommon Mistakes
- Forgetting
@functools.wraps(func)— the decorated function loses its__name__,__doc__, and__module__. Debugging shows “wrapper” instead of the function name. Always add@functools.wraps(func)to every wrapper function inside every decorator. - Forgetting parentheses on parameterized decorators —
@retrywithout parentheses passes the function as the first argument toretry, which expectsmax_attempts. Use@retry()even with all defaults. The parentheses tell Python “call the factory first, then apply the result as a decorator.” - Not returning the result from the wrapper —
def wrapper(*args, **kwargs): func(*args, **kwargs)calls the function but returnsNone. Alwaysreturn func(*args, **kwargs)to pass the return value through. - Returning
Truefrom__exit__accidentally — returningTruesuppresses exceptions. A suppressed exception means your pipeline “succeeds” when it actually failed. Always returnFalse(or nothing) from__exit__unless you specifically intend to swallow the exception. - Not using
yieldin a@contextmanagerfunction — a@contextmanagerfunction must contain exactly oneyield. Without it, Python raises a RuntimeError. With two yields, only the first is used and cleanup code after the second never runs. - Using decorators for side effects on import — a decorator runs when the function is defined (at import time), not when it is called. If your decorator makes database connections or API calls, they happen during
import— before the program even starts. Keep decorators lightweight; defer heavy work to the wrapper function. - Over-decorating — stacking 5+ decorators makes code hard to debug. The stack trace shows wrapper functions instead of actual code. Use decorators for cross-cutting concerns (timing, retry, logging) and keep business logic in the function itself.
Interview Questions
Q: What is a decorator and how does it work?
A: A decorator is a function that takes another function as input and returns a new function that usually extends the original’s behavior. The @decorator syntax is shorthand for func = decorator(func). It works because Python functions are first-class objects that can be passed as arguments and returned from other functions. Decorators use closures to remember the original function and add behavior before or after calling it.
Q: What is a closure? A: A closure is a function that remembers variables from its enclosing scope even after that scope has finished executing. In decorators, the wrapper function is a closure that remembers the original function. Closures work because Python functions carry a reference to their enclosing scope’s variables. This is what allows a decorator’s wrapper to call the original function — it “remembers” it via the closure.
Q: Why is @functools.wraps important?
A: Without @wraps, the decorated function’s __name__, __doc__, and __module__ are replaced with the wrapper’s. This breaks debugging (stack traces show “wrapper” instead of the function name), documentation generation, and serialization frameworks. @wraps(func) copies the original function’s metadata to the wrapper, preserving its identity.
Q: What is the difference between a function-based and class-based decorator?
A: A function-based decorator uses a nested function (closure) to wrap the target. A class-based decorator uses __init__ to receive the function and __call__ to execute it. Class-based decorators are useful when you need to maintain state across calls (like a call counter) because the state lives as instance attributes.
Q: What is a context manager and when do you use one?
A: A context manager is an object with __enter__ and __exit__ methods that the with statement uses for automatic setup and cleanup. __enter__ runs at the start (open file, connect to DB). __exit__ runs at the end — even if an exception occurs (close file, close connection, rollback transaction). Use context managers whenever you have resources that must be released: files, database connections, network sockets, locks, temporary files.
Q: What is contextlib.contextmanager and how does it simplify context managers?
A: @contextlib.contextmanager lets you write a context manager as a generator function with a single yield, instead of a full class with __enter__ and __exit__. Code before yield is the setup (__enter__). The yielded value becomes the as variable. Code after yield (in finally) is the cleanup (__exit__). It is simpler and more readable for straightforward context managers.
Q: How would you build a retry decorator for a production data pipeline?
A: Build a decorator factory that accepts max_attempts, delay, backoff multiplier, and a tuple of exception types to catch. The wrapper loops up to max_attempts, calling the function each time. On failure, it waits delay * backoff^(attempt-1) seconds (exponential backoff) and retries. On the final attempt, it re-raises the exception. Add logging for each retry and failure. Always use @functools.wraps to preserve the function’s identity.
Wrapping Up
Decorators and context managers are the two patterns that separate script-quality Python from production-quality Python. Decorators let you add retry, timing, logging, validation, and caching to any function without touching its code — apply once, reuse everywhere. Context managers guarantee that resources are always cleaned up — connections closed, files deleted, transactions committed or rolled back — no matter what happens inside the block.
The key decorator patterns for data engineering are @timer (performance monitoring), @retry (resilience), @log_call (observability), and @lru_cache (performance). The key context manager patterns are database connections (commit/rollback/close), temporary directories (auto-cleanup), and pipeline audit loggers (structured tracking). Combine them, and your pipeline functions contain only business logic — all the infrastructure concerns are handled by reusable decorators and context managers.
Previous in this series: Regular Expressions — Pattern Matching and Data Cleaning
Next in this series: Working with APIs and HTTP Requests
Naveen Vuppula is a Senior Data Engineering Consultant and app developer based in Ontario, Canada. He writes about Python, SQL, AWS, Azure, and everything data engineering at DriveDataScience.com.