Imagine you are working on a Python application that interacts with a database. Every time you connect to the database, you want to ensure the connection is properly opened and closed, regardless of whether an error occurs during your operations. Rather than manually opening and closing the connection each time, you decide to use a context manager. By implementing a class-based context manager, as you have seen in previous lessons, you can automate the setup and teardown of the simulated database connection, making your code more robust and easier to maintain.


import unittest
import user_code
import ast
import re   
import unittest
import importlib
import io
from contextlib import redirect_stdout

class TestTask(unittest.TestCase):
    def test_enter_prints_connecting(self):
        import user_code
        importlib.reload(user_code)
        f = io.StringIO()
        with redirect_stdout(f):
            with user_code.DatabaseConnection("TestDB") as conn:
                pass
        output = normalize_text(f.getvalue())
        expected = normalize_text("Connecting to database: TestDB")
        _dynamic_test(
            self,
            expected in output,
            f"Output includes connecting message: {expected}",
            f"Expected connecting message: '{expected}' not found in output: '{output}'",
        )

    def test_exit_prints_disconnecting(self):
        import user_code
        importlib.reload(user_code)
        f = io.StringIO()
        with redirect_stdout(f):
            with user_code.DatabaseConnection("TestDB") as conn:
                pass
        output = normalize_text(f.getvalue())
        expected = normalize_text("Disconnecting from database: TestDB")
        _dynamic_test(
            self,
            expected in output,
            f"Output includes disconnecting message: {expected}",
            f"Expected disconnecting message: '{expected}' not found in output: '{output}'",
        )

    def test_supports_different_db_names(self):
        import user_code
        importlib.reload(user_code)
        db_name = "MyDatabase"
        f = io.StringIO()
        with redirect_stdout(f):
            with user_code.DatabaseConnection(db_name) as conn:
                pass
        output = normalize_text(f.getvalue())
        expected_connect = normalize_text(f"Connecting to database: {db_name}")
        expected_disconnect = normalize_text(f"Disconnecting from database: {db_name}")
        _dynamic_test(
            self,
            expected_connect in output and expected_disconnect in output,
            f"Output includes correct messages for arbitrary db_name: {db_name}",
            f"Expected both connecting and disconnecting messages for '{db_name}' 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()
    for i, node in enumerate(tree.body):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == var_name:
                    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] = []
                    
                    return '\\n'.join(lines)
    return code

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


test_main.py

Master the art of resource management in Python by learning how to use and create context managers. This course covers the underlying protocol, practical applications, and advanced techniques for robust and efficient code.

Explore the fundamentals of context managers, their protocol, and how they simplify resource management in Python.

Delve into creating your own context managers using both classes and generator functions.

Solidify your understanding by tackling practical tasks that require context managers in various scenarios.

Challenge: Class-based Resource Manager

Solution

Awesome!

Challenge: Class-based Resource Manager

Solution