Skip to content

For the past five years, you've honed your skills as a Senior Data Scientist for a global university. Your team leverages its data analytics and machine learning skill sets to help other departments make data-driven decisions. One such department is the procurement team, who is trying to decide the best new mobile phone to offer to the university's employees. For the last week, a Junior Data Scientist on your team has been developing a workflow to help provide insight to the procurement team. You will be reviewing their code to ensure it's ready to ship to production.

The first chunk of code that you'll be reviewing is your colleague's function to prepare smartphone data from a CSV file for visualization. After ingesting and cleaning the smartphone data, your colleague has prepared a function to plot a variable passed to the function, versus "price". However, within this function, there is code that does not adhere to DRY principles and is copied and pasted. Make sure to refactor the code appropriately, using the column_to_label() function defined below.

Wow, your colleague even included a unit test to ensure NaN values were removed from the cleaned DataFrame! However, it doesn't seem like the unit test is passing when executed. Re-work this unit test to ensure that it matches the transformation logic in the prepare_smartphone_data() function.

Once you've made changes to the test_nan_values unit test, you'll want to ensure that these unit tests execute with ExitCode.OK. This means that the pytest defined above has passed testing, and the code is one step closer to being to be shipped to production.

For context, there is a print statement in the prepare_smartphone_data() function in the first cell of the notebook below that can be used to visualize the dataset your Junior Data Engineer has been working with. Feel free to update this line of code as needed. This can then be removed after the dataset has been investigated. Best of luck!

import os
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

def prepare_smartphone_data(file_path):
    """
    Prepares smartphone data for visualization by performing the following:
        - Retains only a subset of relevant columns
        - Removes records with missing values for 'battery_capacity' or 'os'
        - Converts the 'price' from cents to dollars

    Parameters:
        file_path (str): Path to the CSV file containing smartphone data

    Returns:
        pd.DataFrame: Cleaned smartphone dataset
    """

    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found at path: {file_path}")

    # Load raw data
    raw_data = pd.read_csv(file_path)

    # Subset columns needed for analysis
    columns_to_keep = [
        "brand_name",
        "os",
        "price",
        "avg_rating",
        "processor_speed",
        "battery_capacity",
        "screen_size"
    ]
    trimmed_data = raw_data[columns_to_keep]

    # Remove records missing critical information
    reduced_data = trimmed_data.dropna(subset=["battery_capacity", "os"])

    # Convert price from cents to dollars
    reduced_data["price"] = reduced_data["price"] / 100

    return reduced_data


# Call the function (ensure this path is correct in your environment)
cleaned_data = prepare_smartphone_data("./data/smartphones.csv")

import seaborn as sns
import matplotlib.pyplot as plt


def column_to_label(column_name):
    """
    Converts a column name in a pandas DataFrame to a string that can be
    used as a label in a plot.
    
    :param column_name: string containing original column name
    :return: string that is ready to be presented on a plot
    """
    if isinstance(column_name, str):
        return " ".join(column_name.split("_")).title()
    else:
        raise Exception("Please make sure to pass a value of type 'str'.")


def visualize_versus_price(clean_data, x):
    """
    Creates a scatter plot to visualize the relationship between a given variable and price.
    
    :param clean_data: A pandas DataFrame containing cleaned smartphone data
    :param x: Column name to be plotted on the x-axis
    :return: None
    """
    
    # Create the scatterplot
    sns.scatterplot(x=x, y="price", data=clean_data, hue="os")
    
    # Use column_to_label() for clean, readable labels
    x_label = column_to_label(x)
    y_label = column_to_label("price")

    # Add x and y labels
    plt.xlabel(x_label)
    plt.ylabel(y_label + " ($)")

    # Add a title to the plot
    plt.title(f"{x_label} vs. {y_label}")
    plt.show()


# Call the function with cleaned data
visualize_versus_price(cleaned_data, "processor_speed")

# Import required packages
import pytest
import ipytest

ipytest.config.rewrite_asserts = True
__file__ = "notebook.ipynb"


# Create a clean DataFrame fixture
@pytest.fixture()
def clean_smartphone_data():
    return prepare_smartphone_data("./data/smartphones.csv")


def test_nan_values(clean_smartphone_data):
    """
    Test for no NaN values in 'battery_capacity' or 'os'
    """
    assert clean_smartphone_data["battery_capacity"].isnull().sum() == 0
    assert clean_smartphone_data["os"].isnull().sum() == 0


ipytest.run("-qq")