Passed
Pull Request — master (#226)
by Steve
02:53
created

AsyncContextContainer.magic_partial()   A

Complexity

Conditions 2

Size

Total Lines 17
Code Lines 13

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 4
CRAP Score 2.032

Importance

Changes 0
Metric Value
cc 2
eloc 13
nop 6
dl 0
loc 17
ccs 4
cts 5
cp 0.8
crap 2.032
rs 9.75
c 0
b 0
f 0
1 1
import inspect
2 1
import logging
3 1
from asyncio import Lock
4 1
from contextlib import AsyncExitStack
5 1
from copy import copy
6 1
from functools import wraps
7 1
from typing import (
8
    Optional,
9
    Type,
10
    TypeVar,
11
    Awaitable,
12
    Generic,
13
    Collection,
14
    Union,
15
    ContextManager,
16
    AsyncContextManager,
17
    Iterator,
18
    Generator,
19
    AsyncGenerator,
20
    Callable,
21
    List,
22
    Dict,
23
)
24
25 1
from lagom.container import Container
26 1
from lagom.definitions import Alias, ConstructionWithContainer, SingletonWrapper
27 1
from lagom.exceptions import InvalidDependencyDefinition, MissingFeature
28 1
from lagom.experimental.definitions import AsyncConstructionWithContainer
29 1
from lagom.interfaces import (
30
    ReadableContainer,
31
    SpecialDepDefinition,
32
    CallTimeContainerUpdate,
33
    ContainerBoundFunction,
34
)
35
36 1
T = TypeVar("T")
37 1
X = TypeVar("X")
38
39
40 1
class AwaitableSingleton(Generic[T]):
41 1
    instance: Optional[T]
42 1
    constructor: ConstructionWithContainer[Awaitable[T]]
43 1
    container: Container
44 1
    _lock: Lock
45
46 1
    def __init__(self, constructor: ConstructionWithContainer, container: Container):
47 1
        self.instance = None
48 1
        self.constructor = constructor  # type: ignore
49 1
        self.container = container
50 1
        self._lock = Lock()
51
52 1
    async def get(self) -> T:
53 1
        if not self.instance:
54 1
            async with self._lock:
55 1
                if not self.instance:
56 1
                    self.instance = await self.constructor.get_instance(self.container)
57 1
        return self.instance
58
59
60 1 View Code Duplication
class _AsyncContextBoundFunction(ContainerBoundFunction[X]):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
61
    """
62
    Represents an instance of a function bound to an async context container
63
    """
64
65 1
    __slots__ = ("async_context_container", "partially_bound_function")
66
67 1
    async_context_container: "AsyncContextContainer"
68 1
    partially_bound_function: ContainerBoundFunction
69
70 1
    __dict__: Dict = dict()
71
72 1
    def __init__(
73
        self,
74
        async_context_container: "AsyncContextContainer",
75
        partially_bound_function: ContainerBoundFunction,
76
    ):
77 1
        self.async_context_container = async_context_container
78 1
        self.partially_bound_function = partially_bound_function
79
80 1
    def __call__(self, *args, **kwargs) -> X:
81 1
        return self.__async_call__(*args, **kwargs)
82
83 1
    async def __async_call__(self, *args, **kwargs):
84 1
        async with self.async_context_container as c:
85 1
            return await self.partially_bound_function.rebind(c)(*args, **kwargs)
86
87 1
    def rebind(self, container: ReadableContainer) -> "ContainerBoundFunction[X]":
88
        return wraps(self.partially_bound_function)(
89
            _AsyncContextBoundFunction(
90
                self.async_context_container,
91
                self.partially_bound_function.rebind(container),
92
            )
93
        )
94
95 1
    def __getattr__(self, item):
96
        if item not in self.__slots__:
97
            raise Exception(f"{item} doesn't exist")
98
        if item == "async_context_container":
99
            return self.async_context_container
100
        if item == "partially_bound_function":
101
            return self.partially_bound_function
102
103
104 1
class AsyncContextContainer(Container):
105 1
    async_exit_stack: Optional[AsyncExitStack] = None
106 1
    _context_types: Collection[Type]
107 1
    _context_singletons: Collection[Type]
108 1
    _root_context: bool = True
109
110 1
    def __init__(
111
        self,
112
        container: Container,
113
        context_types: Collection[Type],
114
        context_singletons: Collection[Type] = tuple(),
115
        log_undefined_deps: Union[bool, logging.Logger] = False,
116
    ):
117 1
        super().__init__(container, log_undefined_deps)
118 1
        self._context_types = set(context_types)
119 1
        self._context_singletons = set(context_singletons)
120
121 1
    def clone(self) -> "AsyncContextContainer":
122
        """returns a copy of the container
123
        :return:
124
        """
125 1
        return AsyncContextContainer(
126
            self,
127
            context_types=self._context_types,
128
            context_singletons=self._context_singletons,
129
            log_undefined_deps=self._undefined_logger,
130
        )
131
132 1
    async def __aenter__(self):
133 1
        if not self.async_exit_stack and self._root_context:
134 1
            self.async_exit_stack = AsyncExitStack()
135
136 1
        if self.async_exit_stack and self._root_context:
137
            # All actual context definitions happen on a clone so that there's isolation between invocations
138 1
            in_context = self.clone()
139 1
            in_context.async_exit_stack = AsyncExitStack()
140 1
            in_context._root_context = False
141
142 1
            for dep_type in self._context_types:
143 1
                managed_dep = self._context_type_def(dep_type)
144 1
                key = Awaitable[dep_type] if isinstance(managed_dep, AsyncConstructionWithContainer) else dep_type  # type: ignore
145 1
                in_context[key] = managed_dep  # type: ignore
146 1
            for dep_type in self._context_singletons:
147 1
                managed_singleton = self._singleton_type_def(dep_type)
148 1
                key = AwaitableSingleton[dep_type] if isinstance(managed_singleton, AwaitableSingleton) else dep_type  # type: ignore
149 1
                in_context[key] = managed_singleton  # type: ignore
150
151
            # The parent context manager keeps track of the inner clone
152 1
            await self.async_exit_stack.enter_async_context(in_context)
153 1
            return in_context
154 1
        return self
155
156 1
    async def __aexit__(self, exc_type, exc_val, exc_tb):
157 1
        if self.async_exit_stack:
158 1
            await self.async_exit_stack.aclose()
159 1
            self.async_exit_stack = None
160
161 1
    def partial(
162
        self,
163
        func: Callable[..., X],
164
        shared: Optional[List[Type]] = None,
165
        container_updater: Optional[CallTimeContainerUpdate] = None,
166
    ) -> ContainerBoundFunction[X]:
167 1
        if not inspect.iscoroutinefunction(func):
168
            raise MissingFeature(
169
                "AsyncContextManager currently can only deal with async functions"
170
            )
171 1
        base_partial = super(AsyncContextContainer, self).partial(
172
            func, shared, container_updater
173
        )
174
175 1
        return wraps(base_partial)(_AsyncContextBoundFunction(self, base_partial))
176
177 1
    def magic_partial(
178
        self,
179
        func: Callable[..., X],
180
        shared: Optional[List[Type]] = None,
181
        keys_to_skip: Optional[List[str]] = None,
182
        skip_pos_up_to: int = 0,
183
        container_updater: Optional[CallTimeContainerUpdate] = None,
184
    ) -> ContainerBoundFunction[X]:
185 1
        if not inspect.iscoroutinefunction(func):
186
            raise MissingFeature(
187
                "AsyncContextManager currently can only deal with async functions"
188
            )
189 1
        base_partial = super(AsyncContextContainer, self).magic_partial(
190
            func, shared, keys_to_skip, skip_pos_up_to, container_updater
191
        )
192
193 1
        return wraps(base_partial)(_AsyncContextBoundFunction(self, base_partial))
194
195 1
    def _context_type_def(self, dep_type: Type):
196 1
        type_def = self.get_definition(ContextManager[dep_type]) or self.get_definition(Iterator[dep_type]) or self.get_definition(Generator[dep_type, None, None]) or self.get_definition(AsyncGenerator[dep_type, None]) or self.get_definition(AsyncContextManager[dep_type])  # type: ignore
197 1
        if type_def is None:
198 1
            raise InvalidDependencyDefinition(
199
                f"A ContextManager[{dep_type}] should be defined. "
200
                f"This could be an Iterator[{dep_type}] or Generator[{dep_type}, None, None] "
201
                f"with the @contextmanager decorator"
202
            )
203 1
        if isinstance(type_def, Alias):
204
            # Without this we create a definition that points to
205
            # itself.
206 1
            type_def = copy(type_def)
207 1
            type_def.skip_definitions = True
208 1
        if self.get_definition(AsyncGenerator[dep_type, None]) or self.get_definition(AsyncContextManager[dep_type]):  # type: ignore
209 1
            return AsyncConstructionWithContainer(lambda c: self._async_context_resolver(c, type_def))  # type: ignore
210 1
        return ConstructionWithContainer(lambda c: self._context_resolver(c, type_def))  # type: ignore
211
212 1
    def _context_resolver(self, c: ReadableContainer, type_def: SpecialDepDefinition):
213
        """
214
        Takes an existing definition which must be a context manager. Returns
215
        the value of the context manager from __enter__ and then places the
216
        __exit__ in this container's exit stack
217
        """
218 1
        assert self.async_exit_stack, "Types can only be resolved within an async with"
219 1
        context_manager = type_def.get_instance(c)
220 1
        return self.async_exit_stack.enter_context(context_manager)
221
222 1
    def _async_context_resolver(
223
        self, c: ReadableContainer, type_def: SpecialDepDefinition
224
    ):
225
        """
226
        Takes an existing definition which must be a context manager. Returns
227
        the value of the context manager from __aenter__ and then places the
228
        __aexit__ in this container's exit stack
229
        """
230 1
        assert self.async_exit_stack, "Types can only be resolved within an async with"
231 1
        context_manager = type_def.get_instance(c)
232 1
        return self.async_exit_stack.enter_async_context(context_manager)
233
234 1
    def _singleton_type_def(self, dep_type: Type):
235
        """
236
        The same as context_type_def but acts as a singleton within this container
237
        """
238 1
        type_def = self._context_type_def(dep_type)
239 1
        if isinstance(type_def, AsyncConstructionWithContainer):
240 1
            return AwaitableSingleton(type_def, self)
241
        return SingletonWrapper(type_def)
242