1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
|
import inspect
from functools import wraps
from typing import TYPE_CHECKING, TypeVar
from typing_extensions import ParamSpec
from beanie.exceptions import StateManagementIsTurnedOff, StateNotSaved
if TYPE_CHECKING:
from beanie.odm.documents import AnyDocMethod, AsyncDocMethod, DocType
P = ParamSpec("P")
R = TypeVar("R")
def check_if_state_saved(self: "DocType"):
if not self.use_state_management():
raise StateManagementIsTurnedOff(
"State management is turned off for this document"
)
if self._saved_state is None:
raise StateNotSaved("No state was saved")
def saved_state_needed(
f: "AnyDocMethod[DocType, P, R]",
) -> "AnyDocMethod[DocType, P, R]":
@wraps(f)
def sync_wrapper(self: "DocType", *args: P.args, **kwargs: P.kwargs) -> R:
check_if_state_saved(self)
return f(self, *args, **kwargs)
@wraps(f)
async def async_wrapper(
self: "DocType", *args: P.args, **kwargs: P.kwargs
) -> R:
check_if_state_saved(self)
# type ignore because there is no nice/proper way to annotate both sync
# and async case without parametrized TypeVar, which is not supported
return await f(self, *args, **kwargs) # type: ignore[misc]
if inspect.iscoroutinefunction(f):
# type ignore because there is no nice/proper way to annotate both sync
# and async case without parametrized TypeVar, which is not supported
return async_wrapper # type: ignore[return-value]
return sync_wrapper
def check_if_previous_state_saved(self: "DocType"):
if not self.use_state_management():
raise StateManagementIsTurnedOff(
"State management is turned off for this document"
)
if not self.state_management_save_previous():
raise StateManagementIsTurnedOff(
"State management's option to save previous state is turned off for this document"
)
def previous_saved_state_needed(
f: "AnyDocMethod[DocType, P, R]",
) -> "AnyDocMethod[DocType, P, R]":
@wraps(f)
def sync_wrapper(self: "DocType", *args: P.args, **kwargs: P.kwargs) -> R:
check_if_previous_state_saved(self)
return f(self, *args, **kwargs)
@wraps(f)
async def async_wrapper(
self: "DocType", *args: P.args, **kwargs: P.kwargs
) -> R:
check_if_previous_state_saved(self)
# type ignore because there is no nice/proper way to annotate both sync
# and async case without parametrized TypeVar, which is not supported
return await f(self, *args, **kwargs) # type: ignore[misc]
if inspect.iscoroutinefunction(f):
# type ignore because there is no nice/proper way to annotate both sync
# and async case without parametrized TypeVar, which is not supported
return async_wrapper # type: ignore[return-value]
return sync_wrapper
def save_state_after(
f: "AsyncDocMethod[DocType, P, R]",
) -> "AsyncDocMethod[DocType, P, R]":
@wraps(f)
async def wrapper(self: "DocType", *args: P.args, **kwargs: P.kwargs) -> R:
result = await f(self, *args, **kwargs)
self._save_state()
return result
return wrapper
|
