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

class TestTask(unittest.TestCase):
    def setUp(self):
        self.pd = __import__('pandas')
        self.df = self.pd.DataFrame({
            "Asset_A": [0.01, 0.02, -0.005, 0.015, 0.007],
            "Asset_B": [0.008, 0.018, -0.004, 0.013, 0.009],
            "Market": [0.009, 0.019, -0.0045, 0.014, 0.008]
        })

    def test_function_exists_and_returns_array(self):
        import user_code
        importlib.reload(user_code)
        func = getattr(user_code, 'predict_asset_returns', None)
        _dynamic_test(
            self,
            callable(func),
            "Function 'predict_asset_returns' is defined.",
            "Function 'predict_asset_returns' is not defined.",
        )
        result = func(self.df) if callable(func) else None
        _dynamic_test(
            self,
            result is not None and isinstance(result, (np.ndarray, list)),
            "Function returns a numpy array or list.",
            f"Returned value is not a numpy array or list: {result}",
        )

    def test_predictions_are_close_to_expected(self):
        import user_code
        importlib.reload(user_code)
        func = getattr(user_code, 'predict_asset_returns', None)
        result = func(self.df) if callable(func) else None
        # Calculate expected using solution
        from sklearn.linear_model import LinearRegression
        X = self.df[["Asset_B", "Market"]].values
        y = self.df["Asset_A"].values
        model = LinearRegression()
        model.fit(X, y)
        expected = model.predict(X)
        _dynamic_test(
            self,
            result is not None and np.allclose(result, expected, atol=1e-8),
            "Predictions are correct and match LinearRegression output.",
            f"Predictions do not match expected values.\\nExpected: {expected}\\nGot: {result}",
        )

    def test_works_with_different_data(self):
        import user_code
        importlib.reload(user_code)
        func = getattr(user_code, 'predict_asset_returns', None)
        # New test data
        df2 = self.pd.DataFrame({
            "Asset_A": [0.02, 0.03, 0.01, 0.025, 0.012],
            "Asset_B": [0.015, 0.025, 0.005, 0.018, 0.011],
            "Market": [0.012, 0.022, 0.004, 0.017, 0.013]
        })
        result = func(df2) if callable(func) else None
        from sklearn.linear_model import LinearRegression
        X = df2[["Asset_B", "Market"]].values
        y = df2["Asset_A"].values
        model = LinearRegression()
        model.fit(X, y)
        expected = model.predict(X)
        _dynamic_test(
            self,
            result is not None and np.allclose(result, expected, atol=1e-8),
            "Function works with new data and returns correct predictions.",
            f"Predictions do not match for new data.\\nExpected: {expected}\\nGot: {result}",
        )

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: Predict Asset Returns with Linear Regression

Solution