add finalize callback function (deepjavalibrary#572)

engines/python/setup/djl_python/outputs.py

@@ -52,6 +52,8 @@ def __init__(self, code=200, message='OK'):
  self.properties = dict()
  self.content = PairList()
  self.stream_content = None
+ self.finalize_function = None
+ self.finalize_args = None
 
  def __str__(self):
  d = dict()
@@ -144,6 +146,15 @@ def write_utf8(msg, val):
  msg += struct.pack('>h', len(buf))
  msg += buf
 
+ def finalize(self, finalize_function, *args):
+ self.finalize_function = finalize_function
+ self.finalize_args = args
+ return self
+
+ def execute_finalize(self):
+ if self.finalize_function:
+ return self.finalize_function(*self.finalize_args)
+
  def send(self, cl_socket):
  msg = bytearray()
  msg += struct.pack('>h', self.code)

engines/python/setup/djl_python/tests/test_input_output.py

@@ -49,6 +49,14 @@ def test_print_message(self):
  [1., 1.]]])]'''
  self.assertEqual(result, expected)
 
+ def test_finalize(self):
+
+ def finalize_func(a, b, c):
+ return a + b + c
+
+ outputs = Output().finalize(finalize_func, 1, 2, 3)
+ self.assertEqual(6, outputs.execute_finalize())
+
 
 if __name__ == '__main__':
  unittest.main()

engines/python/setup/djl_python_engine.py

@@ -122,6 +122,10 @@ def run_server(self):
 
  outputs.send(cl_socket)
  logging.debug("Outputs is sent to DJL engine.")
+ try:
+ outputs.execute_finalize()
+ except Exception as e:
+ logging.exception(f"Failed on finalize function: {e}")
 
 
 def main():