"""Efficient iteration and generator patterns in Python. This module demonstrates various techniques for efficient data processing, including generators, memory profiling, and the use of built-in iteration tools from the `itertools` module. """ from __future__ import annotations import itertools import random import tracemalloc from collections.abc import Generator, Iterable from dataclasses import dataclass from datetime import date from functools import wraps from typing import Any, TypeVar T = TypeVar("T") # --- Lists vs. Tuples --- print("--- Lists vs. Tuples ---") my_list = list(range(1_000_000)) # This creates a list with 1,000,000 integers in memory. print(f"my_list created with {len(my_list)} elements.") # --- Iterators --- print("\n--- Iterators ---") my_list_iterator = [1, 2, 3] my_iterator = iter(my_list_iterator) print(next(my_iterator)) print(next(my_iterator)) print(next(my_iterator)) # --- Generators --- print("\n--- Generators ---") def number_generator(n: int) -> Generator[int]: """Yield numbers from 0 to n-1.""" yield from range(n) gen = number_generator(1_000_000) print("Generator created.") # --- Example: Flattening a List of Lists --- print("\n--- Example: Flattening a List of Lists ---") trades_cashflows: list[Any] = [ [10, 20, 30], [15, 25], [100, -10, 5], 110, ] def flatten(list_of_lists: Iterable[Any]) -> Generator[Any]: """Flatten an iterable of iterables into a single generator.""" for item in list_of_lists: if isinstance(item, list): yield from item else: yield item all_cashflows_generator = flatten(trades_cashflows) for cf in all_cashflows_generator: print(cf, end=" ") print("\n") # --- Memory Efficiency in Action --- print("--- Memory Efficiency in Action ---") def profile_memory(func: Any) -> Any: # noqa: ANN401 """Decorate a function to profile its memory usage.""" @wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 tracemalloc.start() result = func(*args, **kwargs) current, peak = tracemalloc.get_traced_memory() print(f"Function: {func.__name__}") print( f"Current memory usage is {current / 10**6:.6f}MB; " f"Peak was {peak / 10**6:.6f}MB" ) tracemalloc.stop() return result return wrapper @profile_memory def create_list(n_elements: int) -> list[int]: """Create a list of n numbers.""" return list(range(n_elements)) @profile_memory def create_generator(n_elements: int) -> Generator[int]: """Create a generator of n numbers.""" return (i for i in range(n_elements)) n_val = 1_000_000 print("Profiling memory for list creation...") my_list = create_list(n_val) print("\nProfiling memory for generator creation...") my_generator = create_generator(n_val) @profile_memory def consume_generator[T](gen_to_consume: Iterable[T]) -> list[T]: """Consume a generator and return its elements as a list.""" return list(gen_to_consume) print("\nProfiling memory for generator consumption...") consumed_list = consume_generator(my_generator) print("\n") # --- Essential Iteration Tools: enumerate --- print("--- Essential Iteration Tools: enumerate ---") cashflows = [100, 100, 100, 1100] for period, cf in enumerate(cashflows, 1): print(f"Period {period}: Cashflow = {cf}") print("\n") # --- Essential Iteration Tools: zip --- print("--- Essential Iteration Tools: zip ---") trade_dates = ["2025-11-05", "2025-11-06", "2025-11-07"] notionals = [1_000_000, 2_500_000, 500_000] for trade_date, notional in zip(trade_dates, notionals, strict=False): print(f"On {trade_date}, we traded a notional of {notional:,}") print("\n") # --- Essential Iteration Tools: sorted --- print("--- Essential Iteration Tools: sorted ---") @dataclass class Trade: """Represent a financial trade with an ID, maturity date, and notional.""" trade_id: str maturity: date notional: float trades = [ Trade("T1", date(2026, 12, 31), 10_000_000), Trade("T2", date(2025, 12, 31), 5_000_000), Trade("T3", date(2027, 12, 31), 15_000_000), ] sorted_by_maturity = sorted(trades, key=lambda t: t.maturity) for trade in sorted_by_maturity: print(trade) print("\n") # --- The itertools Module: chain.from_iterable --- print("--- The itertools Module: chain.from_iterable ---") fixed_leg = [50, 50, 50, 50] floating_leg = [50 + random.uniform(-5, 5) for _ in range(6)] # noqa: S311 bond_legs = [fixed_leg, floating_leg] full_swap_leg = itertools.chain.from_iterable(bond_legs) print("Full cashflow stream for the leg:") for cf in full_swap_leg: print(f"{cf:.2f}", end=" ") print("\n") # --- The itertools Module: accumulate (P&L) --- print("--- The itertools Module: accumulate (P&L) ---") daily_pnl = [150, -200, 50, 300, -100] cumulative_pnl = itertools.accumulate(daily_pnl) print(list(cumulative_pnl)) print("\n") # --- The itertools Module: accumulate (Amortization) --- print("--- The itertools Module: accumulate (Amortization) ---") def outstanding_balance(balance: float, payment: float, rate: float) -> float: """Calculate the remaining balance after a payment and interest application.""" return balance * (1 + rate) - payment initial_notional = 1_000_000 interest_rate = 0.01 monthly_payment = 5000 payments = itertools.repeat(monthly_payment) balances = itertools.accumulate( payments, lambda balance, pmt: outstanding_balance(balance, pmt, interest_rate), initial=initial_notional, ) amortization_schedule = itertools.takewhile(lambda balance: balance > 0, balances) for i, balance in enumerate(amortization_schedule): print(f"Month {i + 1}: {balance:,.2f}") print("\n") # --- The itertools Module: pairwise --- print("--- The itertools Module: pairwise ---") payment_dates = [date(2025, 1, 15), date(2025, 7, 15), date(2026, 1, 15)] def year_fraction(start_date: date, end_date: date) -> float: """Calculate the year fraction between two dates.""" return (end_date - start_date).days / 365.25 for start, end in itertools.pairwise(payment_dates): yf = year_fraction(start, end) print(f"Period: {start} to {end}, Year Fraction: {yf:.4f}") print("\n") # --- The itertools Module: cycle --- print("--- The itertools Module: cycle ---") scenarios = itertools.cycle(["Base", "Rate Up", "Rate Down"]) for _ in range(5): print(f"Running scenario: {next(scenarios)}")