Building Typed Python Functions and Unit Tests

Organizing Code with Functions

So far, our code has been a single sequence of instructions. As our data processing logic gets more complex, this becomes hard to read and reuse. The solution is to bundle related lines of code into a function. A function is a named, reusable block of code that performs a specific task.

Let's refactor the data processing logic from our previous lesson. Instead of having all the logic in the main script, we can define a function that takes the raw data as input and returns the processed summary.

The article "Defining Your Own Python Function" from Real Python provides an excellent overview of the syntax and concepts.

As you've just read, a function has a name, a list of parameters (inputs), and a body of code. It can optionally return a value (the output).

Here's how we could define a function for our benchmark summarization task:

def summarize_benchmark_data(records: list[dict]) -> list[dict]:
    """
    Calculates the average latency for each model configuration from a list of records.

    Args:
        records: A list of dictionaries, where each dictionary represents a benchmark run.
                 Expected keys: 'model', 'batch_size', 'latency_ms'.

    Returns:
        A list of dictionaries, each summarizing the results for one configuration.
    """
    processed_data = {}
    
    for row in records:
        try:
            model_name = row['model']
            batch_size = int(row['batch_size'])
            latency = float(row['latency_ms'])
            
            config_key = (model_name, batch_size)
            
            if config_key not in processed_data:
                processed_data[config_key] = []
            processed_data[config_key].append(latency)
        except (ValueError, KeyError):



            # We silently skip malformed rows in this version
            pass

    summary_list = []
    for config, latencies in processed_data.items():
        if not latencies:
            continue
        avg_latency = sum(latencies) / len(latencies)
        summary_list.append({
            "model": config[0],
            "batch_size": config[1],
            "num_runs": len(latencies),
            "avg_latency_ms": round(avg_latency, 2)
        })
        
    return summary_list

Notice a few things here:

• The function has a clear name that describes its purpose.
• It takes one argument, records, which is expected to be a list of dictionaries.
• It returns a new list of dictionaries, summary_list, without modifying the original input data. This is a "pure" function with no side effects, which is highly desirable.
• It includes a docstring (the triple-quoted string) that explains what the function does, its arguments (Args), and what it returns (Returns). This is crucial for documentation.
• It uses type hints like records: list[dict] and -> list[dict]. Let's dive into those.

Making Code Clearer with Type Hints

The line def summarize_benchmark_data(records: list[dict]) -> list[dict]: uses type hints. These annotations don't change how the code runs, but they provide valuable information to other developers (and your future self!) about what kind of data the function expects and returns.

• records: list[dict] means the records parameter is expected to be a list containing dictionaries.
• -> list[dict] means the function is expected to return a list of dictionaries.

Type hints make your code more self-documenting, help catch bugs before you even run the code (with tools called static type checkers), and enable better autocompletion in code editors. This is a standard practice in professional Python development.

This video from Corey Schafer gives a great practical walkthrough of using type hints, from the basics to more advanced patterns.

Verifying Behavior with Unit Tests

We've written a function, but how do we know it works correctly? We could run our script with some sample data and check the output manually, but this is tedious and not scalable. A much better approach is to write unit tests: small, automated tests that verify a specific piece of code (a "unit," like our function) behaves as expected.

The introduction to this video explains beautifully why we can't just trust that our code works.

While that video uses Python's built-in unittest module, we will use pytest, a more modern, powerful, and widely-used testing framework that you are likely to encounter in a professional setting. pytest makes writing tests simpler and more intuitive.

The official pytest documentation provides a great starting point.

Writing Tests for Our Function

Let's apply these principles to test our summarize_benchmark_data function. According to pytest conventions, we'll create a new file named test_data_processing.py.

First, you'll need to install pytest:
pip install pytest

Now, here is what test_data_processing.py could look like. It would live in the same directory as the file containing our function.

# test_data_processing.py
import pytest



# Assume summarize_benchmark_data is in a file named data_processing.py
from data_processing import summarize_benchmark_data

def test_summarize_happy_path():
    """Tests the basic functionality with clean, valid data."""
    raw_data = [
        {"model": "Llama-3-8B", "batch_size": "8", "latency_ms": "100.0"},
        {"model": "Mistral-7B", "batch_size": "16", "latency_ms": "50.0"},
        {"model": "Llama-3-8B", "batch_size": "8", "latency_ms": "102.0"},
    ]
    
    summary = summarize_benchmark_data(raw_data)
    



    # Expected output should be sorted to ensure a consistent comparison
    expected = [
        {
            "model": "Llama-3-8B",
            "batch_size": 8,
            "num_runs": 2,
            "avg_latency_ms": 101.0
        },
        {
            "model": "Mistral-7B",
            "batch_size": 16,
            "num_runs": 1,
            "avg_latency_ms": 50.0
        },
    ]




    # Sort both lists by model name to make the comparison order-independent
    summary.sort(key=lambda x: x['model'])
    expected.sort(key=lambda x: x['model'])

    assert summary == expected

def test_summarize_with_malformed_data():
    """Tests that the function skips rows with invalid data."""
    raw_data = [
        {"model": "Llama-3-8B", "batch_size": "8", "latency_ms": "100.0"},
        {"model": "Mistral-7B", "batch_size": "16", "latency_ms": "invalid"},
        {"model": "Gemma-2B", "latency_ms": "25.0"}, # Missing batch_size
    ]
    
    summary = summarize_benchmark_data(raw_data)
    
    expected = [
        {
            "model": "Llama-3-8B",
            "batch_size": 8,
            "num_runs": 1,
            "avg_latency_ms": 100.0
        }
    ]
    
    assert summary == expected

def test_summarize_empty_input():
    """Tests that the function returns an empty list for empty input."""
    summary = summarize_benchmark_data([])
    assert summary == []

def test_summarize_float_precision():
    """Tests floating point comparison using pytest.approx."""
    raw_data = [
        {"model": "Test-Model", "batch_size": "1", "latency_ms": "0.1"},
        {"model": "Test-Model", "batch_size": "1", "latency_ms": "0.2"},
    ]
    summary = summarize_benchmark_data(raw_data)
    



    # 0.1 + 0.2 is not exactly 0.3 in floating point arithmetic
    expected_avg_latency = 0.3 / 2 
    
    assert len(summary) == 1
    assert summary[0]["avg_latency_ms"] == pytest.approx(expected_avg_latency)

To run these tests, you simply navigate to the directory in your terminal and run the command pytest. pytest will automatically find and execute the test functions, reporting which ones passed and which failed. This automated feedback loop is incredibly powerful for ensuring your code remains correct as you add features or refactor it.

Conclusion

In this lesson, you've made a significant leap from writing simple scripts to adopting professional software engineering practices. By packaging your logic into functions, you make it reusable, more readable, and easier to manage.

Here are the key takeaways:

• Functions (def) are the primary way to structure and reuse code in Python.
• Well-designed functions often take data as input and return a result, avoiding "side effects" that modify external state.
• Type hints (:) are annotations that clarify the expected data types for function arguments and return values, improving code clarity and safety.
• Unit tests (using frameworks like pytest) provide an automated way to verify your function's correctness for a wide range of inputs and edge cases.

The code we've written today—a well-defined function with a docstring, type hints, and a separate file of unit tests—is far more robust and maintainable than the script we started with.

In the next lesson, we will continue this journey of code organization. You'll learn how to group related functions and classes into modules, allowing you to build even larger and more structured Python applications.