
import unittest
import user_code
import ast
import re   
import importlib
import csv
import unittest
import importlib
import sys
import io
import numpy as np

class TestTask(unittest.TestCase):
    def setUp(self):
        # Patch plt.show to not block during tests
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        self._orig_show = plt.show
        plt.show = lambda: None

    def tearDown(self):
        # Restore plt.show
        import matplotlib.pyplot as plt
        plt.show = self._orig_show

    def test_corr_matrix_calculation(self):
        import user_code
        importlib.reload(user_code)
        corr_matrix = getattr(user_code, 'corr_matrix', None)
        _dynamic_test(
            self,
            corr_matrix is not None and hasattr(corr_matrix, 'shape') and corr_matrix.shape == (5,5),
            "Correlation matrix is calculated and has correct shape (5,5)",
            f"Correlation matrix not found or shape is not (5,5): {getattr(corr_matrix, 'shape', None)}",
        )

    def test_heatmap_display(self):
        # This test checks that a heatmap is created using seaborn
        import user_code
        import matplotlib.pyplot as plt
        import seaborn as sns
        importlib.reload(user_code)
        # Check if a figure was created
        figs = [fig for fig in plt.get_fignums()]
        _dynamic_test(
            self,
            len(figs) > 0,
            "Heatmap figure was created",
            "No matplotlib figure was created for the heatmap",
        )

    def test_highest_and_lowest_correlation_pair_printed(self):
        import user_code
        importlib.reload(user_code)
        f = io.StringIO()
        sys_stdout = sys.stdout
        sys.stdout = f
        importlib.reload(user_code)
        sys.stdout = sys_stdout
        output = normalize_text(f.getvalue())
        expected_prefixes = [
            "highest correlation:",
            "lowest correlation:"
        ]
        found_prefixes = [prefix for prefix in expected_prefixes if prefix in output]
        _dynamic_test(
            self,
            len(found_prefixes) == 2,
            "Both highest and lowest correlation pairs are printed in the correct format",
            f"Expected both 'Highest correlation:' and 'Lowest correlation:' lines in output, got: {output}",
        )

    def test_highest_and_lowest_pair_names(self):
        import user_code
        importlib.reload(user_code)

        import pandas as pd
        returns = pd.DataFrame({
            'AAPL': [0.01, 0.02, -0.01, 0.03, 0.00, 0.01, 0.02, -0.02, 0.01, 0.02],
            'MSFT': [0.02, 0.01, 0.00, 0.02, 0.01, 0.00, 0.01, -0.01, 0.02, 0.01],
            'GOOG': [0.03, 0.02, -0.02, 0.04, 0.01, 0.02, 0.03, -0.01, 0.02, 0.03],
            'AMZN': [-0.01, 0.01, 0.00, 0.01, -0.02, 0.00, 0.01, -0.01, 0.00, 0.01],
            'TSLA': [0.05, 0.03, -0.04, 0.06, 0.02, 0.03, 0.05, -0.03, 0.04, 0.05]
        })

        corr_matrix = returns.corr()
        corr_pairs = corr_matrix.where(~np.eye(corr_matrix.shape[0], dtype=bool)).stack()
        highest = corr_pairs.idxmax()
        lowest = corr_pairs.idxmin()

        # normalize expected to lowercase because output_norm is lowercase
        highest_pair = {highest[0].lower(), highest[1].lower()}
        lowest_pair = {lowest[0].lower(), lowest[1].lower()}

        f = io.StringIO()
        sys_stdout = sys.stdout
        sys.stdout = f
        importlib.reload(user_code)
        sys.stdout = sys_stdout

        output = normalize_text(f.getvalue())

        found_highest = ("highest correlation:" in output) and all(s in output for s in highest_pair)
        found_lowest = ("lowest correlation:" in output) and all(s in output for s in lowest_pair)

        _dynamic_test(
            self,
            found_highest and found_lowest,
            "Correct pair names for highest and lowest correlation are printed",
            f"Expected highest pair {highest_pair} and lowest pair {lowest_pair} in output: {output}",
        )


def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)

def normalize_text(text):
    text = text.lower()
    text = re.sub(r"\\s{2,}", " ", text)
    text = re.sub(r"\\s*([,:?])\\s*", r"\\1 ", text)
    return text.strip()

def change_var(code: str, var_name: str, value: str) -> str:
    tree = ast.parse(code)
    lines = code.splitlines()
    changed = False
    # Collect all assignment nodes to modify
    assign_nodes = [
        (i, node)
        for i, node in enumerate(tree.body)
        if isinstance(node, ast.Assign)
        and any(isinstance(target, ast.Name) and target.id == var_name for target in node.targets)
    ]

    # If nothing to change, return unmodified code
    if not assign_nodes:
        return code

    # Perform replacements for all matching assignments (from last to first to not break line offsets)
    for i, node in reversed(assign_nodes):
        start_line = node.lineno - 1
        line = lines[start_line]
        indent = ' ' * (len(line) - len(line.lstrip()))
        lines[start_line] = f"{indent}{var_name} = {value}"
        next_line = len(lines)
        for next_node in tree.body[i+1:]:
            if hasattr(next_node, 'lineno'):
                next_line = next_node.lineno - 1
                break
        if next_line > start_line + 1:
            lines[start_line+1:next_line] = []
        changed = True

    return '\\n'.join(lines) if changed else code

if __name__ == "__main__":
    unittest.main()


test_main.py

Unlock the power of Python to analyze financial data, evaluate investment opportunities, and automate key investor workflows. This course guides you through practical, real-world scenarios using Python, focusing on data analysis, visualization, and quantitative techniques relevant to investors.

Learn how to use Python to analyze and visualize financial data, focusing on the core skills investors need to interpret market trends and make informed decisions.

Dive deeper into key investment metrics and learn how to analyze and visualize portfolios using Python.

Apply advanced Python techniques to automate investment analysis and leverage statistical and machine learning tools for deeper insights.

Challenge: Analyze Asset Correlations

Solution