Контекстні менеджери

import time

class Timer:
    def __enter__(self):
        self._start = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.time() - self._start
        return False  # не пригнічуємо виключення


with Timer() as t:
    total = sum(range(1_000_000))

print(f'Elapsed: {t.elapsed:.4f}s')

class ManagedFile:
    def __init__(self, path, mode='r'):
        self.path = path
        self.mode = mode
        self._file = None

    def __enter__(self):
        self._file = open(self.path, self.mode)
        return self._file

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._file:
            self._file.close()
        return False


with ManagedFile('/tmp/test.txt', 'w') as f:
    f.write('привіт контекстний менеджер\n')

with ManagedFile('/tmp/test.txt', 'r') as f:
    print(f.read())

import time
from contextlib import contextmanager

@contextmanager
def timer():
    info = {'elapsed': None}
    start = time.time()
    try:
        yield info           # 'as' отримує цей об'єкт
    finally:
        info['elapsed'] = time.time() - start


with timer() as t:
    total = sum(range(1_000_000))

print(f"Elapsed: {t['elapsed']:.4f}s")

from contextlib import contextmanager

@contextmanager
def suppress_and_log(*exception_types):
    try:
        yield
    except exception_types as e:
        print(f'Suppressed: {e}')


with suppress_and_log(ValueError, ZeroDivisionError):
    result = 1 / 0
    print('цей рядок не буде досягнуто')

print('виконання продовжується після блоку with')
# Suppressed: division by zero
# виконання продовжується після блоку with

import os
import tempfile
from contextlib import contextmanager

@contextmanager
def temp_file(suffix='.txt'):
    path = tempfile.mktemp(suffix=suffix)
    try:
        yield path
    finally:
        if os.path.exists(path):
            os.remove(path)


with temp_file() as path:
    with open(path, 'w') as f:
        f.write('тимчасові дані')
    print(os.path.exists(path))  # True

print(os.path.exists(path))  # False

import builtins
from contextlib import contextmanager

@contextmanager
def indented(level=1, indent='  '):
    prefix = indent * level
    original_print = builtins.print
    def indented_print(*args, **kwargs):
        original_print(prefix, end='')
        original_print(*args, **kwargs)
    builtins.print = indented_print
    try:
        yield
    finally:
        builtins.print = original_print


print('верхній рівень')
with indented(2):
    print('відступ 4 пробіли')
    print('теж відступ')
print('повернення до верхнього рівня')

from contextlib import contextmanager

class FakeDB:
    def __init__(self):
        self.log = []

@contextmanager
def transaction(db):
    pending = []
    try:
        yield pending
        db.log.extend(pending)  # підтвердження при успіху
    except Exception:
        pass  # відкочування — відкидаємо pending
        raise


db = FakeDB()

with transaction(db) as tx:
    tx.append('INSERT user')
    tx.append('UPDATE balance')

print(db.log)  # ['INSERT user', 'UPDATE balance']

try:
    with transaction(db) as tx:
        tx.append('DELETE everything')
        raise RuntimeError('oops')
except RuntimeError:
    pass

print(db.log)  # ['INSERT user', 'UPDATE balance']

import contextlib

def read_first_lines(paths):
    with contextlib.ExitStack() as stack:
        files = [stack.enter_context(open(p)) for p in paths]
        return [f.readline().strip() for f in files]


import tempfile, os

paths = []
for i in range(3):
    p = tempfile.mktemp()
    with open(p, 'w') as f:
        f.write(f'рядок з файлу {i}\nще рядки')
    paths.append(p)

print(read_first_lines(paths))

for p in paths:
    os.remove(p)

class Attempt:
    def __init__(self):
        self.error = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is ValueError:
            self.error = exc_val
            return True  # пригнічуємо виключення
        return False     # інші виключення поширюються


with Attempt() as a:
    int('не число')

print(a.error)

with Attempt() as a:
    x = int('42')

print(a.error)  # None

import threading

class SafeCounter:
    def __init__(self):
        self.count = 0
        self.lock = threading.Lock()

    def increment(self):
        with self.lock:
            self.count += 1


counter = SafeCounter()
threads = [threading.Thread(target=lambda: [counter.increment() for _ in range(1000)])
           for _ in range(5)]
for t in threads: t.start()
for t in threads: t.join()
print(counter.count)  # 5000