|
3 | 3 | Licensed under MIT. |
4 | 4 | ''' |
5 | 5 |
|
6 | | -import imp |
7 | 6 | import sys |
8 | 7 | import traceback |
9 | 8 | import json |
@@ -39,10 +38,32 @@ def filter(self, record): |
39 | 38 | return True |
40 | 39 |
|
41 | 40 |
|
| 41 | +class FunctionLoader(): |
| 42 | + def __init__(self, |
| 43 | + request_id=None, |
| 44 | + source=None, |
| 45 | + function_name=None, |
| 46 | + library_path=None, |
| 47 | + func=None): |
| 48 | + self.request_id = request_id |
| 49 | + self.source = source |
| 50 | + self.function_name = function_name |
| 51 | + self.library_path = library_path |
| 52 | + |
| 53 | + self.func = func |
| 54 | + |
| 55 | + def load(self): |
| 56 | + if self.library_path is not None: |
| 57 | + load_lib(self.library_path) |
| 58 | + |
| 59 | + self.func = load_source( |
| 60 | + self.request_id, self.source, self.function_name) |
| 61 | + |
| 62 | + |
42 | 63 | def call(func, event, context, environment_variables={}): |
43 | 64 | export_variables(environment_variables) |
44 | | - |
45 | | - return _runner(func, event, context) |
| 65 | + loader = FunctionLoader(func=func) |
| 66 | + return _runner(loader, event, context) |
46 | 67 |
|
47 | 68 |
|
48 | 69 | def run(args): |
@@ -54,17 +75,19 @@ def run(args): |
54 | 75 | args.timeout, |
55 | 76 | invoked_function_arn=args.arn_string, |
56 | 77 | function_version=args.version_name) |
57 | | - if args.library is not None: |
58 | | - load_lib(args.library) |
59 | | - func = load(c.aws_request_id, args.file, args.function) |
| 78 | + loader = FunctionLoader( |
| 79 | + request_id=c.aws_request_id, |
| 80 | + source=args.file, |
| 81 | + function_name=args.function, |
| 82 | + library_path=args.library) |
60 | 83 |
|
61 | | - (result, err_type) = _runner(func, e, c) |
| 84 | + (result, err_type) = _runner(loader, e, c) |
62 | 85 |
|
63 | 86 | if err_type is not None: |
64 | 87 | sys.exit(EXITCODE_ERR) |
65 | 88 |
|
66 | 89 |
|
67 | | -def _runner(func, event, context): |
| 90 | +def _runner(loader, event, context): |
68 | 91 | logger = logging.getLogger() |
69 | 92 |
|
70 | 93 | logger.info("Event: {}".format(event)) |
@@ -74,7 +97,7 @@ def _runner(func, event, context): |
74 | 97 | queue = multiprocessing.Queue() |
75 | 98 | p = multiprocessing.Process( |
76 | 99 | target=execute_in_process, |
77 | | - args=(queue, func, event, context,)) |
| 100 | + args=(queue, loader, event, context,)) |
78 | 101 | p.start() |
79 | 102 | (result, err_type, duration) = queue.get() |
80 | 103 | p.join() |
@@ -95,14 +118,25 @@ def load_lib(path): |
95 | 118 | sys.path.append(os.path.abspath(path)) |
96 | 119 |
|
97 | 120 |
|
98 | | -def load(request_id, path, function_name): |
| 121 | +def load_source(request_id, path, function_name): |
99 | 122 | mod_name = 'request-' + str(request_id) |
100 | 123 |
|
101 | 124 | file_path = os.path.abspath(path) |
102 | 125 | file_directory = os.path.dirname(file_path) |
103 | 126 | sys.path.append(file_directory) |
104 | 127 |
|
105 | | - mod = imp.load_source(mod_name, path) |
| 128 | + if sys.version_info.major == 2: |
| 129 | + import imp |
| 130 | + mod = imp.load_source(mod_name, path) |
| 131 | + elif sys.version_info.major == 3 and sys.version_info.minor >= 5: |
| 132 | + import importlib |
| 133 | + spec = importlib.util.spec_from_file_location(mod_name, path) |
| 134 | + mod = importlib.util.module_from_spec(spec) |
| 135 | + sys.modules[mod_name] = mod |
| 136 | + spec.loader.exec_module(mod) |
| 137 | + else: |
| 138 | + raise Exception("unsupported python version") |
| 139 | + |
106 | 140 | func = getattr(mod, function_name) |
107 | 141 | return func |
108 | 142 |
|
@@ -132,9 +166,11 @@ def execute(func, event, context): |
132 | 166 | return result, err_type |
133 | 167 |
|
134 | 168 |
|
135 | | -def execute_in_process(queue, func, event, context): |
| 169 | +def execute_in_process(queue, loader, event, context): |
| 170 | + if loader.func is None: |
| 171 | + loader.load() |
136 | 172 | start_time = timeit.default_timer() |
137 | | - result, err_type = execute(func, event, context) |
| 173 | + result, err_type = execute(loader.func, event, context) |
138 | 174 | end_time = timeit.default_timer() |
139 | 175 | duration = (end_time - start_time) * 1000 |
140 | 176 |
|
|