AskAiml
Sign in Get started

Optimisation techniques for pure Python

Avoiding attribute lookups, caching, __slots__, and when NumPy or Cython is the right escape hatch.

⚡ Module 7 10 min read Not started

Why this matters

Once profiling has identified the hot path, you have a menu of optimisations: algorithmic fixes (always start here), better data structures, NumPy vectorisation, Cython/Numba/mypyc compilation, calling out to Rust. This lesson walks the ladder from "rewrite the loop" to "ship a Rust extension."

Learning objectives

  1. Pick the right optimisation level for the problem.
  2. Apply algorithmic and data-structure fixes first.
  3. Vectorise with NumPy / Polars.
  4. JIT with Numba or compile with Cython / mypyc.
  5. Bind to Rust with PyO3 (overview).

1. The optimisation ladder

Pure Python (idiomatic)            ──  fast enough for most code
   ↓
Algorithm change (O(n²) → O(n))    ──  free 10-1000x; do this first
   ↓
Better data structure              ──  set vs list, dict, deque, heapq
   ↓
NumPy / Polars vectorisation       ──  10-100x for numeric / tabular work
   ↓
Numba JIT (@njit)                  ──  drop-in for numeric loops; near C speed
   ↓
Cython / mypyc                     ──  ahead-of-time C compilation
   ↓
PyO3 / Rust / C extension          ──  ultimate speed; biggest engineering cost

Climb only as high as needed. Almost no one needs PyO3.

2. Step 1: algorithmic fixes

Membership: list → set

python
# BAD — O(n²)
keepers = [x for x in items if x in valid_ids]   # valid_ids: list

# GOOD — O(n)
valid_ids = set(valid_ids)
keepers = [x for x in items if x in valid_ids]

Look-up: list scan → dict

python
# BAD — O(n × m)
matched = [next(u for u in users if u.id == oid) for oid in order_ids]

# GOOD — O(n + m)
by_id = {u.id: u for u in users}
matched = [by_id[oid] for oid in order_ids]

Repeated computation → cache

python
@functools.cache
def expensive(x: int) -> int: ...

Repeated regex → compile once

python
PATTERN = re.compile(r"...")
PATTERN.search(s)                   # vs re.search at every call

Linear search → bisect

python
import bisect
i = bisect.bisect_left(sorted_list, x)

O(log n) for sorted arrays.

Priority queue → heapq

python
import heapq
heap = []
heapq.heappush(heap, (priority, item))
priority, item = heapq.heappop(heap)

3. Step 2: better data structures

deque for FIFO / sliding windows

python
from collections import deque
window = deque(maxlen=100)

list.pop(0) is O(n). deque.popleft() is O(1).

dict over many if/elif

python
HANDLERS = {"a": handle_a, "b": handle_b, "c": handle_c}
HANDLERS.get(kind, handle_default)(payload)

array for homogeneous numeric arrays

python
import array
nums = array.array("i", range(10**7))

10× less memory than a Python list of ints.

slots for mass instances

python
@dataclass(slots=True)
class Point: x: float; y: float

50% less memory per instance.

frozenset for immutable lookups

python
ALLOWED: frozenset[str] = frozenset({"a", "b", "c"})

Hashable, immutable, fast in.

4. Step 3: vectorise with NumPy

(See Lesson 5.1.) Convert loops to array ops.

Before:

python
def haversine(lat1, lon1, lat2, lon2):
    # ... math on scalars ...
    return d

distances = [haversine(*coords[i], *coords[j]) for i, j in pairs]

After:

python
import numpy as np

def haversine_vec(lat1, lon1, lat2, lon2):
    # all inputs are numpy arrays; same math, vector-wise
    ...
    return d

coords = np.array(coords)              # (N, 2)
distances = haversine_vec(
    coords[i_idx, 0], coords[i_idx, 1],
    coords[j_idx, 0], coords[j_idx, 1],
)

10-100× speedup for size > 1000.

Polars over pandas

Already covered (5.3). 3-30× speedup with the same API.

numexpr for big expressions

python
import numexpr as ne
result = ne.evaluate("a*b + c*d - e**2", local_dict={"a": a, "b": b, ...})

Avoids intermediates; uses cache better.

5. Step 4: JIT with Numba

numba compiles Python functions to LLVM at first call. Best for numerical loops where NumPy doesn't naturally express the operation.

powershell
uv add numba
python
from numba import njit
import numpy as np

@njit
def monte_carlo_pi(n: int) -> float:
    inside = 0
    for _ in range(n):
        x = np.random.random()
        y = np.random.random()
        if x*x + y*y <= 1.0:
            inside += 1
    return 4 * inside / n

print(monte_carlo_pi(10_000_000))       # ~100x faster than pure Python

Caveats

  • First call compiles → slow. Use @njit(cache=True) to persist.
  • Only supports a subset of Python (mostly numeric).
  • No classes, no most stdlib calls.
  • For arrays, use NumPy first; Numba shines for loops that NumPy can't express.

@njit(parallel=True) + prange

python
from numba import njit, prange
@njit(parallel=True)
def f(a):
    for i in prange(len(a)):
        a[i] = a[i] * 2

Multi-threaded loops. Bypasses the GIL inside @njit.

6. Step 5: Cython — ahead-of-time C

For ~maximal speed on numeric code, with full Python access:

cython
# myext.pyx
def add(int a, int b) -> int:
    return a + b

cdef double sum_squares(double[::1] arr):
    cdef double s = 0
    cdef int i
    for i in range(arr.shape[0]):
        s += arr[i] * arr[i]
    return s

def sum_squares_py(double[::1] arr):
    return sum_squares(arr)

Build with cython + a build system (setuptools / scikit-build / hatchling).

toml
# pyproject.toml
[build-system]
requires = ["setuptools", "Cython", "numpy"]

[tool.setuptools]
ext-modules = [{name = "myext", sources = ["myext.pyx"]}]
powershell
uv pip install -e .                     # builds the .so / .pyd

Cython gives 10-100× over pure Python on tight loops, with full control over memory layout, types, GIL release (with nogil:). Worth the build complexity for libraries; rarely for app code.

mypyc — compile typed Python

powershell
uv add mypy[mypyc]
mypyc src/myproject

mypyc reads your type hints and generates C code. Much less code change than Cython — just type your existing Python. 2-10× speedup typical.

mypy itself uses mypyc to compile itself. Same for some pydantic-v2 components.

7. Step 6: Rust extensions with PyO3

For ultimate performance and CPython interop:

rust
// src/lib.rs
use pyo3::prelude::*;

#[pyfunction]
fn sum_squares(xs: Vec<f64>) -> f64 {
    xs.iter().map(|x| x * x).sum()
}

#[pymodule]
fn my_native(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_squares, m)?)?;
    Ok(())
}

Build with maturin:

powershell
uv tool install maturin
maturin develop                          # builds + installs into the venv

pip install from PyPI works for users (pre-built wheels). They don't need a Rust toolchain.

Examples in 2026: pydantic-core (Rust core), ruff, uv, polars are all Rust-backed Python libs.

When to consider:

  • You ship a library with hot inner loops used by many.
  • You need to release the GIL and use raw threads.
  • You need memory layout / lifetime guarantees Python doesn't offer.

Not for app code unless the inner loop genuinely matters.

8. Step 7: don't write Python here

Sometimes the right answer is "use the right tool":

  • ML training → PyTorch / JAX (delegates to CUDA).
  • Image processing → OpenCV / Pillow-SIMD.
  • Linear algebra → NumPy with MKL/Accelerate.
  • Database aggregations → DuckDB / Postgres.
  • Compression → zstandard, lz4, xxhash (C-backed).
  • Tokenizers → tokenizers (Rust).
  • Async networking → httpx / aiohttp (C event loop via libuv).

These libraries are 1000× faster than reimplementing the algorithm. Choose the right primitive first.

9. Common Python micro-optimisations (when nothing else helps)

SlowFast
for i in range(len(a)):for x in a: or enumerate(a)
a.append(x) in a loop, then ''.join(a)precompute, or use io.StringIO
Dict lookup of attributelocal-binding (f = self.method)
Global lookup in tight loopbind to local (g = some_func)
repr(x) in hot loggingguard with log.isEnabledFor(DEBUG)
Many small dict allocationsreuse one dict, .clear() between iterations
class X: pass for tiny records__slots__ or NamedTuple
try/except on the hot pathcheck the condition with if first if cheaper
+ to concatenate many strings"".join(parts)

Each is worth ~10-30% in tight loops. Don't reach for these before algorithmic fixes.

10. Caching strategies

In-memory: functools.cache / lru_cache

python
@cache
def slow(x): ...

For pure functions whose arguments are hashable.

Disk: diskcache, joblib.Memory

python
from diskcache import Cache
cache = Cache("./cache")
@cache.memoize(expire=3600)
def fetch(url): ...

For expensive results that survive process restart.

Remote: Redis

python
import redis
r = redis.Redis()
key = f"feature:{user_id}"
if (v := r.get(key)) is None:
    v = compute(user_id)
    r.set(key, v, ex=300)

For sharing cache across processes / machines.

Application-level memoisation table

For "compute and cache N items at once":

python
def get_many(ids):
    missing = [i for i in ids if i not in CACHE]
    if missing:
        fresh = bulk_fetch(missing)
        CACHE.update(fresh)
    return {i: CACHE[i] for i in ids}

Avoids the N+1 cost of single-key cache lookups.

11. Worked example: from 60s to 600ms

Suppose process_orders(orders) takes 60s for 1M orders.

  1. Profile (py-spy): top function is find_user(uid) (linear list scan).
  2. Build users_by_id = {u.id: u for u in users} once. Down to 8s.
  3. Profile again: now compute_total(order) is hot — does a regex on each row.
  4. Compile regex once, hoist out of loop. Down to 4s.
  5. Profile: many small dict allocations. Switch order objects to dataclass(slots=True). Down to 2s.
  6. The remaining hot loop is purely numeric (sums per group). Switch to NumPy / Polars groupby. Down to 0.6s.
  7. Stop. 100× faster, total work: half a day.

You'd be tempted to start at step 6. Profile first — steps 1-5 were easier and proportionally larger wins.

Hands-on lab (2 hours)

  1. Take a slow script (your own or a recipe). Profile it. List the top 3 hotspots.
  2. Replace one list in membership check with set in. Re-measure.
  3. Vectorise one numeric loop with NumPy. Re-measure.
  4. Apply @functools.cache to one pure function with hashable args. Re-measure.
  5. Use Numba @njit on one numeric function; compare first-call vs warm.
  6. (Optional) Compile a small module with mypyc; measure speedup.
  7. (Optional) Write a tiny PyO3 function (string parsing, math); install with maturin develop.

Common pitfalls

  1. Optimising without profiling.
  2. Rewriting in NumPy when the bottleneck is I/O.
  3. Numba on functions that hit pure-Python code → slow fallback.
  4. Caching impure functions → wrong results.
  5. Cython without # cython: language_level=3 (or # distutils: language = c++) → ancient defaults.
  6. PyO3 / Rust for an app feature that runs 10 times a day.

Self-check

  1. State the optimisation ladder.
  2. Why is set in faster than list in?
  3. When use Numba over NumPy?
  4. What does mypyc do?
  5. State one library written in Rust that Python users rely on.

References

Sign in to save your progress and earn badges.