Add support for choosing IPython shell

flapi/__main__.py

@@ -6,20 +6,24 @@
 from argparse import ArgumentParser
 from .cli import install_main, shell_main, uninstall_main
 from pathlib import Path
+from typing import Optional, Any
 
 cli = ArgumentParser(
     description='CLI tool for the Flapi library',
 )
 subparsers = cli.add_subparsers(dest="subcommand")
 
 
-def subcommand(args=[], parent=subparsers):
+def subcommand(args: Optional[list[Any]] = None, parent=subparsers):
     """
     A neat little decorator for reducing my immense confusion when dealing with
     the Argparse library
 
     Source: https://mike.depalatis.net/blog/simplifying-argparse.html
     """
+    if args is None:
+        args = []
+
     def decorator(func):
         parser = parent.add_parser(func.__name__, description=func.__doc__)
         for arg in args:
@@ -74,10 +78,18 @@ def uninstall(args):
     uninstall_main(args.data_dir, args.yes)
 
 
-@subcommand()
+@subcommand([
+    argument(
+        "-s",
+        "--shell",
+        type=str,
+        help="The shell to use with Flapi. Either 'ipython' or 'python'.",
+        default=None,
+    )
+])
 def shell(args):
     """Launch a shell connected to FL Studio"""
-    shell_main()
+    shell_main(args.shell)
 
 
 def main():

flapi/cli/shell.py

@@ -5,11 +5,53 @@
 """
 import sys
 import code
+from typing import Optional
 from flapi import enable, init, disable, heartbeat, fl_exec, fl_eval
 from flapi import __consts as consts
+try:
+    import IPython
+    from IPython import start_ipython
+    from traitlets.config.loader import Config as IPythonConfig
+except ImportError:
+    IPython = None
+    start_ipython = None
+    IPythonConfig = None
 
 
-def shell_main():
+SHELL_SCOPE = {
+    "enable": enable,
+    "init": init,
+    "disable": disable,
+    "heartbeat": heartbeat,
+    "fl_exec": fl_exec,
+    "fl_eval": fl_eval,
+}
+
+
+def start_python_shell():
+    """
+    Start up Python's built-in shell
+    """
+    code.interact(
+        banner="",
+        local=SHELL_SCOPE,
+    )
+
+
+def start_ipython_shell():
+    """
+    Start up an Ipython shell
+    """
+    assert IPython is not None
+    assert start_ipython is not None
+    assert IPythonConfig is not None
+    config = IPythonConfig()
+    config.TerminalInteractiveShell.banner1 \
+        = f"IPython version: {IPython.__version__}"
+    start_ipython(argv=[], user_ns=SHELL_SCOPE, config=config)
+
+
+def shell_main(shell_to_use: Optional[str] = None):
     """Main function to set up the Python shell"""
     print("Flapi interactive shell")
     print(f"Client version: {'.'.join(str(n) for n in consts.VERSION)}")
@@ -31,14 +73,17 @@ def shell_main():
     print("Imported functions:")
     print("enable, init, disable, heartbeat, fl_exec, fl_eval")
 
-    code.interact(
-        banner='',
-        local={
-            'enable': enable,
-            'init': init,
-            'disable': disable,
-            'heartbeat': heartbeat,
-            'fl_exec': fl_exec,
-            'fl_eval': fl_eval,
-        }
-    )
+    if shell_to_use == "python":
+        return start_python_shell()
+    elif shell_to_use == "ipython":
+        if IPython is None:
+            print("Error: IPython is not installed!")
+            exit(1)
+        return start_ipython_shell()
+    else:
+        # Default: launch IPython if possible, but fall back to the default
+        # shell
+        if IPython is not None:
+            return start_ipython_shell()
+        else:
+            return start_python_shell()