Python for Data Engineering
Intermediate 14 min read Python for Data Engineering

Functions and Modules

Shape your code so it ships, not just runs

Well-shaped functions, type hints, dataclasses and module boundaries are what make a pipeline readable, testable and importable — the difference between a script and a system.

functionstype hintsdataclassesmodulespackaging
Prerequisites: Python Basics Review
Before we start

What you will learn

  • Design small, testable functions with clear boundaries
  • Type-hint signatures and model records with dataclasses
  • Structure a pipeline as importable modules instead of one script
Introduction

What is this?

Functions convert a script into reusable logic; modules give it import boundaries; dataclasses give records explicit shape. In production data code these three turn a 'run once' script into something a team can test, review and operate for years.

Motivation

Why do you need it?

  • Small functions each cover one behavior, which is exactly what pytest checks.
  • Type hints and dataclasses make contracts visible — a schema for your code, not just your data.
  • Module boundaries let you import one transform in a notebook or test instead of re-running an entire script.
Think about it

Mental model

Imagine

A workshop's labeled tool drawers

Think of

Modules and functions

You don't keep every tool in one pile; each drawer holds a labelled set. Modules group related functions so the right drawer can be pulled out and the right tool run in isolation.

Mechanics

How it works

  • A function isolates a unit of behavior with explicit inputs/outputs; purity (no hidden global state) makes it testable.
  • Type hints (`def transform(df: pd.DataFrame) -> pd.DataFrame`) document and (via mypy) check contracts.
  • `@dataclass` defines record types with fields and equality, replacing ad-hoc dicts.
  • Modules are `.py` files; packages are directories with `__init__.py`, imported by dotted path.
  • A pipeline package exposes entry points (a `main()`) that a DAG or CLI calls.

From script to importable package.

One big script
Functions
Modules
Types + dataclasses
Testable package
Hands-on

Examples

A typed, testable extraction function

A typed, testable extraction function
from dataclasses import dataclass
from typing import Iterator
 
@dataclass(frozen=True)
class Event:
    event_id: str
    amount: float
    ts: str
 
def parse_events(lines: Iterator[str]) -> Iterator[Event]:
    for line in lines:
        eid, amount, ts = line.strip().split(",")
        yield Event(eid, float(amount), ts)

The parser is a pure function over an iterator — trivially unit tested without files or I/O.

Real world

In production

  • An extraction module exposes `extract_orders(source) -> DataFrame` and is imported by both the DAG and a notebook.
  • Dataclasses model config and record shapes, so mypy catches a mis-typed field at CI time instead of 3am.
  • A transforms package releases semantic versions; dags pin the version they were written against.
Guidelines

Best practices

  • Keep functions pure (inputs in, return out) so tests don't need global state or mocks of the real world.
  • Return small dataclasses or DataFrames rather than throwing ragged dicts around.
  • Name modules by responsibility (extract/transform/load/quality), not by tool.
  • Add a `main()` guarded by `if __name__ == "__main__":` so the module is both importable and runnable.
Watch out

Common mistakes

  • A 400-line function that 'does the pipeline' — impossible to unit test and to blame.
  • Global mutable state (a shared list of processed ids) that makes runs non-reproducible.
  • Modules that import each other in a cycle, breaking imports at 3am.
  • Type hints that lie (`Any` everywhere) — worse than no hints because tooling trusts them.
Root cause analysis

Common issues

Error
ImportError: cannot import name
Root cause

Circular imports or a renamed symbol not picked up

Fix

Break the cycle by importing inside functions or reordering module-level imports.

Error
Two team members get different bundles of the same data from one pipeline
Root cause

Shared mutable state accumulated across runs in the same process

Fix

Make the pipeline pure per run: fresh containers per invocation, no module-level caches.

When things break

Debugging tips

  • When a pipeline misbehaves, run the single suspected function in isolation with real inputs.
  • Enable mypy in CI and treat it as a gate — moving contracts earlier catches field bugs.
  • Put the smallest repro in a pytest file; if it passes, the bug is in composition, not the unit.
Interviews

Interview checkpoint

Test yourself

Quick quiz

0/2
1

Which property makes a function easy to unit test?

2

Type hints are enforced at runtime by default in Python.

Go deeper

Official documentation

Type hints — PEP 484
docs
dataclasses — Python docs
docs