<commit_before>import asyncio
from typing import TypeVar
from collections.abc import Awaitable
import logging

from .typing import AsyncSource, AsyncSink
from .futures import Subscription, chain_future

log = logging.getLogger(__name__)

T = TypeVar("T")


class SubscriptionFactory(Awaitable):
    """A helper class that makes it possible to setup subscriptions both
    using await and async-with."""

    def __init__(self, source, sink=None):
        self._source = source
        self._sink = sink

        self._subscription = None

    async def create(self) -> Subscription:
        """Awaits subscription.

        Awaits until subscription has been setup, and returns
        a subscription future. The future will resolve to the
        last value received in the stream of values."""

        if self._sink is not None:
            down = await Subscription().__alisten__(self._sink)
        else:
            down = Subscription()

        up = await self._source.__alisten__(down)
        self._subscription = chain_future(down, up)

        return self._subscription

    async def __aenter__(self) -> Subscription:
        """Awaits subscription."""
        return await self.create()

    async def __aexit__(self, type, value, traceback) -> None:
        """Closes subscription."""
        self._subscription.cancel()

    def __await__(self) -> Subscription:
        """Await subscription."""
        return self.create().__await__()


def listen(source: AsyncSource, sink: AsyncSink=None) -> SubscriptionFactory:
    return SubscriptionFactory(source, sink)


async def run(source: AsyncSource, asink: AsyncSink, timeout: int=2) -> T:
    """Awaits until subscription closes and returns the final value"""
    return await asyncio.wait_for(await listen(source, asink), timeout)<commit_msg>Sink argument is optional for run()<commit_after>import asyncio
from typing import TypeVar, Optional
from collections.abc import Awaitable
import logging

from .typing import AsyncSource, AsyncSink
from .futures import Subscription, chain_future

log = logging.getLogger(__name__)

T = TypeVar("T")


class SubscriptionFactory(Awaitable):
    """A helper class that makes it possible to setup subscriptions both
    using await and async-with."""

    def __init__(self, source, sink=None):
        self._source = source
        self._sink = sink

        self._subscription = None

    async def create(self) -> Subscription:
        """Awaits subscription.

        Awaits until subscription has been setup, and returns
        a subscription future. The future will resolve to the
        last value received in the stream of values."""

        if self._sink is not None:
            down = await Subscription().__alisten__(self._sink)
        else:
            down = Subscription()

        up = await self._source.__alisten__(down)
        self._subscription = chain_future(down, up)

        return self._subscription

    async def __aenter__(self) -> Subscription:
        """Awaits subscription."""
        return await self.create()

    async def __aexit__(self, type, value, traceback) -> None:
        """Closes subscription."""
        self._subscription.cancel()

    def __await__(self) -> Subscription:
        """Await subscription."""
        return self.create().__await__()


def listen(source: AsyncSource, sink: Optional[AsyncSink]=None) -> SubscriptionFactory:
    return SubscriptionFactory(source, sink)


async def run(source: AsyncSource[T], asink: Optional[AsyncSink]=None, timeout: int=2) -> T:
    """Awaits until subscription closes and returns the final value"""
    return await asyncio.wait_for(await listen(source, asink), timeout)
{'fraction_non_alphanumeric': 0.09352909189777053, 'fraction_numerical': 0.000543773790103317, 'mean_word_length': 3.1759364358683313, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2069443', 'n_tokens_mistral': 973, 'n_tokens_neox': 964, 'n_words': 329}