1313import warnings # to warn if pymongo is missing
1414from contextlib import suppress
1515from datetime import datetime
16+ from typing import Any , Optional , Tuple
1617
1718from .._types import HashFunc , Mongetter
19+ from ..config import CacheEntry
1820
1921with suppress (ImportError ):
2022 from bson .binary import Binary # to save binary data to mongodb
@@ -65,29 +67,29 @@ def __init__(
6567 def _func_str (self ) -> str :
6668 return _get_func_str (self .func )
6769
68- def get_entry_by_key (self , key ) :
70+ def get_entry_by_key (self , key : str ) -> Tuple [ str , Optional [ CacheEntry ]] :
6971 res = self .mongo_collection .find_one (
7072 {"func" : self ._func_str , "key" : key }
7173 )
7274 if not res :
7375 return key , None
7476 try :
75- entry = {
76- " value" : pickle .loads (res ["value" ]), # noqa: S301
77- " time" : res .get ("time" , None ),
78- " stale" : res .get ("stale" , False ),
79- " being_calculated" : res .get ("being_calculated" , False ),
80- }
77+ entry = CacheEntry (
78+ value = pickle .loads (res ["value" ]), # noqa: S301
79+ time = res .get ("time" , None ),
80+ stale = res .get ("stale" , False ),
81+ being_calculated = res .get ("being_calculated" , False ),
82+ )
8183 except KeyError :
82- entry = {
83- " value" : None ,
84- " time" : res .get ("time" , None ),
85- " stale" : res .get ("stale" , False ),
86- " being_calculated" : res .get ("being_calculated" , False ),
87- }
84+ entry = CacheEntry (
85+ value = None ,
86+ time = res .get ("time" , None ),
87+ stale = res .get ("stale" , False ),
88+ being_calculated = res .get ("being_calculated" , False ),
89+ )
8890 return key , entry
8991
90- def set_entry (self , key , func_res ) :
92+ def set_entry (self , key : str , func_res : Any ) -> None :
9193 thebytes = pickle .dumps (func_res )
9294 self .mongo_collection .update_one (
9395 filter = {"func" : self ._func_str , "key" : key },
@@ -104,14 +106,14 @@ def set_entry(self, key, func_res):
104106 upsert = True ,
105107 )
106108
107- def mark_entry_being_calculated (self , key ) :
109+ def mark_entry_being_calculated (self , key : str ) -> None :
108110 self .mongo_collection .update_one (
109111 filter = {"func" : self ._func_str , "key" : key },
110112 update = {"$set" : {"being_calculated" : True }},
111113 upsert = True ,
112114 )
113115
114- def mark_entry_not_calculated (self , key ) :
116+ def mark_entry_not_calculated (self , key : str ) -> None :
115117 with suppress (OperationFailure ): # don't care in this case
116118 self .mongo_collection .update_one (
117119 filter = {
@@ -122,22 +124,22 @@ def mark_entry_not_calculated(self, key):
122124 upsert = False , # should not insert in this case
123125 )
124126
125- def wait_on_entry_calc (self , key ) :
127+ def wait_on_entry_calc (self , key : str ) -> Any :
126128 time_spent = 0
127129 while True :
128130 time .sleep (MONGO_SLEEP_DURATION_IN_SEC )
129131 time_spent += MONGO_SLEEP_DURATION_IN_SEC
130132 key , entry = self .get_entry_by_key (key )
131133 if entry is None :
132134 raise RecalculationNeeded ()
133- if not entry [ " being_calculated" ] :
134- return entry [ " value" ]
135+ if not entry . being_calculated :
136+ return entry . value
135137 self .check_calc_timeout (time_spent )
136138
137- def clear_cache (self ):
139+ def clear_cache (self ) -> None :
138140 self .mongo_collection .delete_many (filter = {"func" : self ._func_str })
139141
140- def clear_being_calculated (self ):
142+ def clear_being_calculated (self ) -> None :
141143 self .mongo_collection .update_many (
142144 filter = {
143145 "func" : self ._func_str ,
0 commit comments