Initial WIP version of nuplugin.py

diff --git a/nuplugin.py b/nuplugin.py
new file mode 100755 (executable)
index 0000000..0260857
--- /dev/null
+++ b/nuplugin.py
@@ -0,0 +1,235 @@
+import argparse
+import os
+import socket
+import sys
+import msgpack
+
+class server:
+    ins = None
+    outs = None
+    errs = None
+    got_hello = False
+    shutdown = False
+    response_thread = None
+    unpacker = None
+    packer = None
+    compatible_version = '0.102.0'
+    features = []
+
+    def __init__(self, ins, outs, errs=None, allow_localsocket=True):
+        self.ins = ins
+        self.outs = outs
+        self.errs = errs
+        self.got_hello = False
+        self.shutdown = False
+        self.unpacker = msgpack.Unpacker(self.ins, raw=True)
+        self.packer = msgpack.Packer(self.outs)
+        if allow_localsocket:
+            features.append({'name':'LocalSocket'})
+
+    def __send_encoding(self):
+        self.outs.write(b'\x07msgpack')
+
+    def __send_hello(self):
+        hello_msg = {'protocol':'nu-plugin', 'version':self.compatible_version, 'features':self.features}
+        self.__send(hello_msg)
+
+    def __send(self, msg):
+        self.packer.pack(msg)
+
+    def send_error_response(self, call_id, error_msg):
+        response = {'CallResponse': [call_id, {'Error': {'msg': error_msg}}]}
+        self.enqueue_response(response)
+
+    def handle_custom_value_unsupported(self, call_id, value):
+        # Custom values aren't yet supported by this module
+        error_msg = "Custom values are not yet supported"
+        self.error(error_msg)
+        self.send_error_response(call_id, error_msg)
+
+    def run_custom_value_op(self, call_id, value, op):
+        op_name = op
+        op_arg = None
+        if type(op) is dict:
+            # It should be a dict of just one key, and the key is the name of the operation being performed
+            if len(op) != 1:
+                error_msg = "Unrecognized CustomValueOp format: ({})".format(op)
+                self.error(error_msg)
+                self.send_error_response(call_id, error_msg)
+                return
+
+            op_name = [key for key in op.keys()][0]
+            op_arg = op.get(op_name)
+
+        match op_name:
+            case 'ToBaseValue':
+                self.handle_custom_value_unsupported(call_id, op_arg)
+            case 'FollowPathInt':
+                self.handle_custom_value_unsupported(call_id, op_arg)
+            case 'FollowPathString':
+                self.handle_custom_value_unsupported(call_id, op_arg)
+            case 'PartialCmp':
+                self.handle_custom_value_unsupported(call_id, op_arg)
+            case 'Operation':
+                self.handle_custom_value_unsupported(call_id, op_arg)
+            case 'Dropped':
+                self.handle_custom_value_unsupported(call_id, op_arg)
+            case _:
+                error_msg = 'Unrecognized CustomValue operation: {}'.format(op_name)
+                self.error(error_msg)
+                self.send_error_response(call_id, error_msg)
+
+    def do_hello(self, msg):
+        protocol = msg.get('protocol')
+        version = msg.get('version')
+        ok = False
+
+        if protocol != 'nu-plugin':
+            self.error('Invalid protocol ({}) sent by engine'.format(protocol))
+        elif not version:
+            self.error('No version number provided by engine')
+        elif not self.compatible_version(version):
+            self.error('Version number {} provided by the engine is incompatible with this plugin.'.format(version))
+        else:
+            ok = True
+
+        if not ok:
+            # Abort communication, terminate...
+            self.shutdown = True
+
+    def do_call(self, msg):
+        [call_id, call_data] = msg['Call']
+        call_name = call_data
+        call_arg = None
+        if type(call_data) is 'dict':
+            if (len(call_data) != 1):
+                error_msg = 'Unrecognized call format: ({})'.format(call_data)
+                self.error(error_msg)
+                self.send_error_response(call_id, error_msg)
+                return
+            (call_name, call_arg) = [(k, v) for (k, v) in call_data.items()][0]
+
+        match call_name:
+            case 'Metadata':
+                do_call_metadata(self, call_id)
+            case 'Signature':
+                do_call_signature(self, call_id)
+            case 'Run':
+                self.run_method(call_id, call_data.get('Run'))
+            case 'CustomValueOp':
+                value = call_arg[0]
+                op = call_arg[1]
+                self.run_custom_value_op(call_id, value, op)
+            case _:
+                error_msg = 'Unrecognized Call: {}'.format(call_name)
+                self.error(error_msg)
+                self.send_error_response(call_id, error_msg)
+
+    def receive_loop(self):
+        while not self.shutdown:
+            try:
+                msg = self.unpacker.unpack()
+                msg_ok = False;
+                msg_type = msg
+                msg_data = None
+
+                if type(msg) is dict:
+                    if len(msg) != 1:
+                        self.error('Unsupported message format: ({})'.format(msg))
+                        return
+                    (msg_type, msg_data) = [(k, v) for (k, v) in msg.items()][0]
+
+                match msg_type:
+                    case 'Hello':
+                        self.do_hello(msg_data)
+                    case 'Call':
+                        self.do_call(msg_data)
+                    case 'Ack':
+                        self.do_ack(msg_data)
+                    case 'Drop':
+                        self.do_drop(msg_data)
+                    case 'EngineCallResponse':
+                        self.error('Got engine call response: Not yet supported by this module')
+                    case 'Goodbye':
+                        self.shutdown = True
+                    case _:
+                        self.error('Got unrecognized message from engine: ({})'.format(msg))
+
+            except msgpack.exceptions.OutOfData:
+                self.error('Engine disconnected without \'Goodbye\' message: Shutting down')
+                self.shutdown = True
+
+    def run(self):
+        # Announce plugin to the shell
+        self.__send_encoding()
+        self.__send_hello()
+
+        # Start the response thread:
+        #  Users of this class may opt to handle requests from Nu shell asynchronously and issue responses out-of-order
+        #  Responses are queued as they're provided to the module server, and the response thread synchronizes access to the out stream for issuing the responses
+        self.start_response_thread()
+
+        # Run the message receipt/processing loop
+        self.receive_loop()
+
+        # If the response thread is still going, shut it down
+        self.shutdown = True
+        if self.response_thread:
+            self.response_thread.join()
+
+        return 0
+
+
+def print_usage():
+    msg = f"To load module, call \"plugin add {sys.argv[0]}\" from nu shell\n"
+    msg += "\nArguments:\n\t--stdio: Run plugin with I/O over stdin/stdout channels\n\t--local-socket <filename>: Run plugin and accept connection on the named local socket for I/O\n"
+    print(msg, file=sys.stderr)
+
+def run_stdio_server():
+    outfd = os.dup(sys.stdout.fileno())
+    outs = os.fdopen(outfd, 'wb')
+    infd = os.dup(sys.stdin.fileno())
+    ins = os.fdopen(infd, 'wb')
+
+    module = server(ins, outs, sys.stderr)
+    return module.run()
+
+def run_local_socket_server(socket_path):
+    # Connect to the socket once to get the in-stream (for receiving from Nu Shell)
+    # Connect a second time to get the out-stream (for sending to Nu Shell)
+    # (The idea is to have the sending and receiving channels be independent,
+    #  so for instance during shutdown the read channel can be closed,
+    #  preventing more commands from being sent to the module,
+    #  while the write channel remains open to finish out previously-queued requests)
+    recv_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM);
+    send_socket = socket_socket(socket.AF_UNIX, socket.SOCK_STREAM);
+
+    recv_socket.connect(socket_path)
+    send_socket.connect(socket_path)
+
+    server = nu_module_server(recv_socket, send_socket, sys.stderr)
+    return server.run()
+
+def main(argv):
+    if len(argv) == 0:
+        print_usage()
+        return 1
+
+    arg_parser = argparse.ArgumentParser(prog='test plugin', description='A simple loadable plugin for Nu Shell, implemented in Python')
+    iogroup = arg_parser.add_mutually_exclusive_group()
+    iogroup.add_argument('--stdio', action='store_true')
+    iogroup.add_argument('--local-socket')
+    options = arg_parser.parse_args(argv)
+
+    if options.stdio:
+        print("testmod starting in stdio mode", file=sys.stderr)
+        return run_stdio_server()
+    if options.local_socket:
+        print("testmod starting in local socket mode ({})".format(options.local_socket), file=sys.stderr)
+        return run_local_socket_server(options.local_socket)
+
+    print(options, file=sys.stderr)
+    return 0
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv[1:]))