From bf8b69d33862e4867be6a0e25cf1dee146231722 Mon Sep 17 00:00:00 2001
From: Thomas Chamberlin <tchamber@nrao.edu>
Date: Wed, 15 Nov 2023 17:56:54 -0500
Subject: [PATCH] First draft of dysh shell

---
 pyproject.toml            |  1 +
 src/dysh/fits/__init__.py |  4 +-
 src/dysh/shell.py         | 88 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 src/dysh/shell.py

diff --git a/pyproject.toml b/pyproject.toml
index 86e03555..8d36b189 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -46,6 +46,7 @@ Source = "https://github.com/GreenBankObservatory/dysh"
 #dysh = "dysh.dysh:main"
 [project.scripts]
 fdr = "dysh.util.files:main_cli"
+dysh = "dysh.shell:main"
 
 [tool.hatch.version]
 path = "src/dysh/__init__.py"
diff --git a/src/dysh/fits/__init__.py b/src/dysh/fits/__init__.py
index a5dc4ce7..b3597907 100644
--- a/src/dysh/fits/__init__.py
+++ b/src/dysh/fits/__init__.py
@@ -1,2 +1,4 @@
 """Classes and functions for importing SDFITS files"""
-__all__ = ["gbtfitsload", "sdfitsload", "gb20mfitsload"]
+from dysh.fits.gb20mfitsload import GB20MFITSLoad  # noqa
+from dysh.fits.gbtfitsload import GBTFITSLoad  # noqa
+from dysh.fits.sdfitsload import SDFITSLoad  # noqa
diff --git a/src/dysh/shell.py b/src/dysh/shell.py
new file mode 100644
index 00000000..e39744dc
--- /dev/null
+++ b/src/dysh/shell.py
@@ -0,0 +1,88 @@
+import argparse
+from pathlib import Path
+from typing import List
+
+import IPython
+from traitlets.config import Config
+
+import dysh.fits
+from dysh import __version__
+from dysh.fits.sdfitsload import SDFITSLoad
+
+# TODO: Derive URLs from pyproject.toml?
+BANNER = f"""--------------------------------------------------------------------------
+                         Welcome to Dysh v{__version__}
+
+    Example usage: https://dysh.readthedocs.io/en/latest/example.html
+    Bug repots:    https://github.com/GreenBankObservatory/dysh/issues
+
+    For help with a Dysh routine from the command line,
+    use the builtin 'help'. e.g.:
+
+        help(GBTFITSLoad)
+
+    The following items have been made globally available for convenience:
+{{user_ns_str}}
+
+   To suppress this banner, start with --no-banner flag
+--------------------------------------------------------------------------
+"""
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Dysh interactive shell.\n\n All CLI arguments other than those defined below are passed through "
+            "to ipython; see $ ipython --help for more details"
+        )
+    )
+    parser.add_argument("paths", help="FITS file paths to load initially", nargs="*", type=Path)
+    parser.add_argument("-p", "--profile", help="The IPython profile to use", default="dysh")
+    parser.add_argument("-L", "--fits-loader", help="The SDFITS loader class name to use", default="GBTFITSLoad")
+    parser.add_argument(
+        "--colors", help="Set the color scheme", choices=["NoColor", "Neutral", "Linux", "LightBG"], default="LightBG"
+    )
+    return parser.parse_known_args()
+
+
+def init_shell(colors: str, profile: str | Path, *ipython_args, sdfits_files=None):
+    c = Config()
+    import numpy as np
+    import pandas as pd
+    from astropy.io import fits
+    from astropy.table import Table
+
+    from dysh.fits.gbtfitsload import GBTFITSLoad
+
+    user_ns = {"pd": pd, "np": np, "GBTFITSLoad": GBTFITSLoad, "Table": Table, "fits": fits}
+    if sdfits_files:
+        user_ns["sdfits_files"] = sdfits_files
+
+    c.BaseIPythonApplication.profile = profile
+    c.InteractiveShell.colors = colors
+    c.InteractiveShell.banner2 = BANNER.format(
+        user_ns_str="\n".join(f"{' '*8}{k} (from {v.__name__})" for k, v in user_ns.items())
+    )
+    IPython.start_ipython(ipython_args, config=c, user_ns=user_ns)
+
+
+def get_fits_loader_class(loader_class_name: str):
+    try:
+        return getattr(dysh.fits, loader_class_name)
+    except AttributeError as error:
+        raise NotImplementedError(f"No known SDFITS Loader {loader_class_name!r}") from error
+
+
+def open_sdfits_files(paths: List[Path], loader_class_name="GBTFITSLoad") -> List[SDFITSLoad]:
+    loader_class = get_fits_loader_class(loader_class_name)
+    return [loader_class(path) for path in paths]
+
+
+def main():
+    args, remaining_args = parse_args()
+    sdfits_files = open_sdfits_files(args.paths, args.fits_loader)
+    init_shell(colors=args.colors, profile=args.profile, *remaining_args, sdfits_files=sdfits_files)
+
+
+if __name__ == "__main__":
+    main()