5
5
import tempfile
6
6
import time
7
7
from http import HTTPStatus
8
+ from typing import Any , Callable , TypeVar , cast
8
9
9
10
import filelock
10
11
import flask
12
+ from flask .typing import ResponseReturnValue
11
13
12
14
import processor
13
15
14
-
15
16
# The file will be flock()'ed if a report is being processed.
16
17
LOCK_FILE = '/tmp/results-processor.lock'
17
18
# If the file above is locked, this timestamp file contains the UNIX timestamp
32
33
app = flask .Flask (__name__ )
33
34
34
35
35
- def _atomic_write (path , content ) :
36
+ def _atomic_write (path : str , content : str ) -> None :
36
37
# Do not auto-delete the file because we will move it after closing it.
37
38
temp = tempfile .NamedTemporaryFile (mode = 'wt' , delete = False )
38
39
temp .write (content )
@@ -41,12 +42,15 @@ def _atomic_write(path, content):
41
42
os .replace (temp .name , path )
42
43
43
44
44
- def _serial_task (func ):
45
+ F = TypeVar ('F' , bound = Callable [..., Any ])
46
+
47
+
48
+ def _serial_task (func : F ) -> F :
45
49
lock = filelock .FileLock (LOCK_FILE )
46
50
47
51
# It is important to use wraps() to preserve the original name & docstring.
48
52
@functools .wraps (func )
49
- def decorated_func (* args , ** kwargs ) :
53
+ def decorated_func (* args : object , ** kwargs : object ) -> object :
50
54
try :
51
55
with lock .acquire (timeout = 1 ):
52
56
return func (* args , ** kwargs )
@@ -55,24 +59,24 @@ def decorated_func(*args, **kwargs):
55
59
return ('A result is currently being processed.' ,
56
60
HTTPStatus .SERVICE_UNAVAILABLE )
57
61
58
- return decorated_func
62
+ return cast ( F , decorated_func )
59
63
60
64
61
- def _internal_only (func ) :
65
+ def _internal_only (func : F ) -> F :
62
66
@functools .wraps (func )
63
- def decorated_func (* args , ** kwargs ) :
67
+ def decorated_func (* args : object , ** kwargs : object ) -> object :
64
68
if (not app .debug and
65
69
# This header cannot be set by external requests.
66
70
# https://cloud.google.com/tasks/docs/creating-appengine-handlers?hl=en#reading_app_engine_task_request_headers
67
71
not flask .request .headers .get ('X-AppEngine-QueueName' )):
68
72
return ('External requests not allowed' , HTTPStatus .FORBIDDEN )
69
73
return func (* args , ** kwargs )
70
74
71
- return decorated_func
75
+ return cast ( F , decorated_func )
72
76
73
77
74
78
@app .route ('/_ah/liveness_check' )
75
- def liveness_check ():
79
+ def liveness_check () -> ResponseReturnValue :
76
80
lock = filelock .FileLock (LOCK_FILE )
77
81
try :
78
82
lock .acquire (timeout = 0.1 )
@@ -91,7 +95,7 @@ def liveness_check():
91
95
92
96
93
97
@app .route ('/_ah/readiness_check' )
94
- def readiness_check ():
98
+ def readiness_check () -> ResponseReturnValue :
95
99
lock = filelock .FileLock (LOCK_FILE )
96
100
try :
97
101
lock .acquire (timeout = 0.1 )
@@ -106,7 +110,7 @@ def readiness_check():
106
110
@app .route ('/api/results/process' , methods = ['POST' ])
107
111
@_internal_only
108
112
@_serial_task
109
- def task_handler ():
113
+ def task_handler () -> ResponseReturnValue :
110
114
_atomic_write (TIMESTAMP_FILE , str (time .time ()))
111
115
112
116
task_id = flask .request .headers .get ('X-AppEngine-TaskName' )
0 commit comments