additional context can be passed to the user-defined callback function, more type hinting

Author: eshaan7

examples/custom_save_fn.py

@@ -14,20 +14,22 @@
 ENDPOINT = "echo"
 
 
-def intercept_result(future: Future):
+def intercept_result(context, future: Future):
     """
     Will be invoked on every process completion
     """
     data = None
     if future.done():
-        with open("/path/to/saved/file") as f:
+        fname = context.get("read_result_from", None)
+        with open(fname) as f:
             data = f.read()
         # 1. get current result object
         res = future.result()
         # 2. update the report variable,
         # you may update only these: report,error,status
         res.report = data
-        res.status = "success"
+        if context.get("force_success", False):
+            res.status = "success"
         # 3. set new result
         future._result = res
 
@@ -42,7 +44,13 @@ def intercept_result(future: Future):
     app.testing = True
     c = app.test_client()
     # request new process
-    data = {"args": ["hello", "world"]}
+    data = {
+        "args": ["hello", "world"],
+        "callback_context": {
+            "read_result_from": "/path/to/saved/file",
+            "force_success": True,
+        },
+    }
     r = c.post(f"/{ENDPOINT}", json=data)
     # get result
     result_url = r.get_json()["result_url"]

examples/run_script.py

@@ -25,6 +25,7 @@ def test():
     ```
     """
     url = "http://localhost:4000/scripts/hacktheplanet"
+    # request without any data
     resp = requests.post(url)
     resp_data = resp.json()
     print(resp_data)

examples/with_callback.py

@@ -14,13 +14,14 @@
 ENDPOINT = "echo"
 
 
-def my_callback_fn(future: Future):
+def my_callback_fn(extra_callback_context, future: Future):
     """
     Will be invoked on every process completion
     """
     print("[i] Process running ?:", future.running())
     print("[i] Process completed ?:", future.done())
     print("[+] Result: ", future.result())
+    print("[+] Context: ", extra_callback_context)
 
 
 shell2http.register_command(

examples/with_signals.py

@@ -22,7 +22,7 @@
 
 
 @my_signal.connect
-def my_callback_fn(future: Future):
+def my_callback_fn(extra_callback_context, future: Future):
     """
     Will be invoked on every process completion
     """

flask_shell2http/api.py

@@ -8,7 +8,8 @@
 
 # system imports
 from http import HTTPStatus
-from typing import Callable, Any
+import functools
+from typing import Callable, Dict, Any
 
 # web imports
 from flask import request, jsonify, make_response
@@ -75,7 +76,9 @@ def post(self):
                 f"Requester: '{request.remote_addr}'."
             )
             # Check if request data is correct and parse it
-            cmd, timeout, key = request_parser.parse_req(request, self.command_name)
+            cmd, timeout, callback_context, key = request_parser.parse_req(
+                request, self.command_name
+            )
 
             # run executor job in background
             future = self.executor.submit_stored(
@@ -84,8 +87,10 @@ def post(self):
             # callback that removes the temporary directory
             future.add_done_callback(request_parser.cleanup_temp_dir)
             if self.user_callback_fn:
-                # user defined callback fn
-                future.add_done_callback(self.user_callback_fn)
+                # user defined callback fn with callback_context if any
+                future.add_done_callback(
+                    functools.partial(self.user_callback_fn, callback_context)
+                )
 
             logger.info(f"Job: '{key}' --> added to queue for command: {cmd}")
             result_url = f"{request.base_url}?key={key}"
@@ -101,7 +106,7 @@ def post(self):
     def __init__(
         self,
         command_name: str,
-        user_callback_fn: Callable[[Future], Any],
+        user_callback_fn: Callable[[Dict, Future], Any],
         executor: Executor,
     ):
         self.command_name: str = command_name

flask_shell2http/base_entrypoint.py

@@ -1,6 +1,6 @@
 # system imports
 from collections import OrderedDict
-from typing import Callable, Any
+from typing import Callable, Dict, Any
 
 # web imports
 from flask_executor import Executor
@@ -74,7 +74,7 @@ def register_command(
         self,
         endpoint: str,
         command_name: str,
-        callback_fn: Callable[[Future], Any] = None,
+        callback_fn: Callable[[Dict, Future], Any] = None,
     ) -> None:
         """
         Function to map a shell command to an endpoint.
@@ -89,17 +89,19 @@ def register_command(
                   For example,
                   if you pass ``{ "args": ["Hello", "World"] }``
                   in POST request, it gets converted to ``echo Hello World``.\n
-            callback_fn (func):
+            callback_fn (Callable[[Dict, Future], Any]):
                 - An optional function that is invoked when a requested process
                     to this endpoint completes execution.
                 - This is added as a
                     ``concurrent.Future.add_done_callback(fn=callback_fn)``
-                - A same callback function may be used for multiple commands.
+                - The same callback function may be used for multiple commands.
+                - if request JSON contains a `callback_context` attr, it will be passed
+                  as the first argument to this function.
 
         Examples::
 
-            def my_callback_fn(future):
-                print(future.result())
+            def my_callback_fn(context: dict, future: Future) -> None:
+                print(future.result(), context)
 
             shell2http.register_command(endpoint="echo", command_name="echo")
             shell2http.register_command(
@@ -108,7 +110,7 @@ def my_callback_fn(future):
                 callback_fn=my_callback_fn
             )
         """
-        uri = self.__construct_route(endpoint)
+        uri: str = self.__construct_route(endpoint)
         # make sure the given endpoint is not already registered
         cmd_already_exists = self.__commands.get(uri)
         if cmd_already_exists:

flask_shell2http/classes.py

@@ -139,33 +139,34 @@ def __parse_multipart_req(args: List[str], files) -> (List[str], str):
         logger.debug(f"Request files saved under temp directory: '{tmpdir}'")
         return args, tmpdir
 
-    def parse_req(self, request, base_command: str) -> (str, str):
+    def parse_req(self, request, base_command: str) -> (str, int, Dict, str):
+        # default values if request is w/o any data
+        # i.e. just run-script
         args: List[str] = []
+        timeout: int = DEFAULT_TIMEOUT
         tmpdir = None
+        callback_context = {}
         if request.is_json:
             # request does not contain a file
             args = request.json.get("args", [])
             timeout: int = request.json.get("timeout", DEFAULT_TIMEOUT)
+            callback_context = request.json.get("callback_context", {})
         elif request.files:
             # request contains file
+            callback_context = request.form.get("callback_context", {})
             received_args = request.form.getlist("args")
             timeout: int = request.form.get("timeout", DEFAULT_TIMEOUT)
             args, tmpdir = RequestParser.__parse_multipart_req(
                 received_args, request.files
             )
-        else:
-            # request is w/o any data
-            # i.e. just run-script
-            args = []
-            timeout: int = DEFAULT_TIMEOUT
 
         cmd: List[str] = base_command.split(" ")
         cmd.extend(args)
         key: str = calc_hash(cmd)
         if tmpdir:
             self.__tmpdirs.update({key: tmpdir})
 
-        return cmd, timeout, key
+        return cmd, timeout, callback_context, key
 
     def cleanup_temp_dir(self, future: Future) -> None:
         key: str = future.result().key