1 files changed, 265 insertions, 0 deletions
diff --git a/Omni/Repl.py b/Omni/Repl.py
new file mode 100644
index 0000000..1cf2f65
--- /dev/null
+++ b/Omni/Repl.py
@@ -0,0 +1,265 @@
+"""
+Improve the standard Python REPL.
+
+This module attempts to emulate the workflow of ghci or lisp repls. It uses
+importlib to load a namespace from the provided path, typechecks it with mypy,
+and provides some tools for improving repl-driven development.
+
+This module is called in Omni/Ide/repl.sh like so:
+
+    python -i Omni/Repl.py NS PATH
+
+where NS is the dot-partitioned namespace of the main module, and PATH is the
+path to the same file. In the future this could be expanded to be a list of
+additional files to load.
+"""
+
+# : out python-improved-repl
+# : dep mypy
+import importlib
+import importlib.util
+import inspect
+import logging
+import mypy.api
+import Omni.Log as Log
+import os
+import pathlib
+import pydoc
+import string
+import subprocess
+import sys
+import textwrap
+import types
+import typing
+import unittest
+
+
+class ReplError(Exception):
+    """Type for errors at the repl."""
+
+
+def use(ns: str, path: str) -> None:
+    """
+    Load or reload the module named 'ns' from 'path'.
+
+    Like `use` in the Guile Scheme repl.
+
+    Raises:
+        ReplError: if module cannot be loaded
+    """
+    logging.info("loading %s from %s", ns, path)
+    spec = importlib.util.spec_from_file_location(ns, path)
+    if spec is None or spec.loader is None:
+        msg = f"spec could not be loaded for {ns} at {path}"
+        raise ReplError(msg)
+    module = importlib.util.module_from_spec(spec)
+    # delete module and its imported names if its already loaded
+    if ns in sys.modules:
+        del sys.modules[ns]
+        for name in module.__dict__:
+            if name in globals():
+                del globals()[name]
+    sys.modules[ns] = module
+    spec.loader.exec_module(module)
+    names = list(module.__dict__)
+    globals().update({k: getattr(module, k) for k in names})
+
+
+def typecheck(path: str) -> None:
+    """Typecheck this namespace."""
+    # this envvar is undocumented, but it works
+    # https://github.com/python/mypy/issues/13815
+    os.environ["MYPY_FORCE_COLOR"] = "1"
+    logging.info("typechecking %s", path)
+    stdout, stderr, _ = mypy.api.run([path])
+    sys.stdout.write(stdout)
+    sys.stdout.flush()
+    sys.stderr.write(stderr)
+    sys.stderr.flush()
+
+
+def edit_file(ns: str, path: str, editor: str) -> None:
+    """
+    Edit and reload the given namespace and path.
+
+    It is assumed ns and path go together. If `editor` returns something other
+    than 0, this function will not reload the ns.
+    """
+    try:
+        proc = subprocess.run([editor, path], check=False)
+    except FileNotFoundError:
+        logging.exception("editor '%s' not found", editor)
+    if proc.returncode == 0:
+        use(ns, path)
+        typecheck(path)
+
+
+class CustomRepl:
+    """Custom repl commands, heavily inspired by ghci."""
+
+    def __init__(self, ns: str, path: str, editor: str) -> None:
+        """Create the custom repl for given ns and path."""
+        self.ns = ns
+        self.path = path
+        self.editor = editor
+        self.default = sys.__excepthook__
+        self.isframe = inspect.isframe
+        self.stack = inspect.stack
+        self.stdout = sys.stdout
+        self.whitespace = string.whitespace
+
+    def setup(self) -> None:
+        """
+        Load the NS, setup hooks and prompt.
+
+        This basically does all the heavy lifting of customizing the Python
+        repl.
+        """
+        # load given namespace
+        use(self.ns, self.path)
+        typecheck(self.path)
+        # setup hooks and prompt
+        sys.excepthook = self.excepthook
+        pydoc.pager = lambda text: pydoc.pipepager(text, "more")
+        sys.ps1 = f"{self.ns}> "
+        sys.ps2 = f"{self.ns}| "
+
+    def help(self) -> str:
+        """Return help text."""
+        return textwrap.dedent(f"""
+            repl commands:
+              :e        open {self.ns} in {self.editor}
+              :r        reload {self.ns}
+              :t obj    show the type of obj
+              obj?      expands to 'help(obj)'
+              :?        show this help
+        """)
+
+    def show_help(self) -> None:
+        """Print info about how to use this repl."""
+        sys.stdout.write(self.help())
+        sys.stdout.flush()
+
+    def excepthook(
+        self,
+        type_: type[BaseException],
+        value: BaseException,
+        traceback: types.TracebackType | None,
+    ) -> typing.Any:
+        """
+        Pre-process Python repl exceptions.
+
+        This is called on `sys.excepthook`, which runs when the repl doesn't
+        know how to handle some input. So, we inspect `value` and provide
+        alternate functionality, bottoming out at the default exception.
+
+        Raises:
+            ReplError: if caught exception is empty
+        """
+        # ruff: noqa: PLR0911
+        if not isinstance(value, SyntaxError):
+            return self.default(type_, value, traceback)
+        if value.text is None:
+            msg = f"value.text is None: {value}"
+            raise ReplError(msg)
+        stmt = value.text.rstrip()
+        if stmt == ":?":
+            self.show_help()
+            return None
+        if stmt.endswith("?"):
+            name = stmt.rstrip("?(" + self.whitespace)
+            self.get_help(name)
+            return None
+        if stmt == ":e":
+            self.edit()
+            return None
+        if stmt == ":r":
+            self.reload()
+            return None
+        if stmt.startswith(":t"):
+            var = stmt.split()[1]
+            self.get_type(var)
+            return None
+        return self.default(type_, value, traceback)
+
+    def get_type(self, name: str) -> typing.Any | None:
+        """Return the type of `name` to the caller."""
+        for record in self.stack():
+            frame = record[0]
+            if not self.isframe(frame):
+                continue
+            cmd = f"typing.reveal_type({name})"
+            return eval(cmd, frame.f_globals, frame.f_locals)  # noqa: S307
+        return None
+
+    def get_help(self, name: str) -> typing.Any | None:
+        """Return the documentation for `name` to the caller."""
+        for record in self.stack():
+            frame = record[0]
+            if not self.isframe(frame):
+                continue
+            cmd = f"help({name})"
+            return eval(cmd, frame.f_globals, frame.f_locals)  # noqa: S307
+        return None
+
+    def reload(self) -> None:
+        """Reload the current namespace."""
+        use(self.ns, self.path)
+        typecheck(self.path)
+
+    def edit(self) -> None:
+        """Edit the current namespace."""
+        edit_file(self.ns, self.path, self.editor)
+
+
+class TestCustomRepl(unittest.TestCase):
+    """Test the CustomRepl functionality."""
+
+    def setUp(self) -> None:
+        """Create a CustomRepl for testing."""
+        ns = __name__
+        path = pathlib.Path(__name__.replace(".", "/"))
+        path = path.with_suffix(".py")
+        self.repl = CustomRepl(ns, str(path), "true")
+        self.repl.setup()
+
+    def tearDown(self) -> None:
+        """Undo `self.setUp`."""
+        sys.excepthook = self.repl.default
+        del self.repl
+
+    def test_help(self) -> None:
+        """Help message should include the ns and path."""
+        self.assertIn(self.repl.ns, self.repl.help())
+
+
+def test() -> None:
+    """Run this module's test suite."""
+    suite = unittest.TestSuite()
+    suite.addTests(
+        unittest.defaultTestLoader.loadTestsFromTestCase(TestCustomRepl),
+    )
+    unittest.TextTestRunner().run(suite)
+
+
+def move() -> None:
+    """Actual entrypoint."""
+    Log.setup()
+    ns = sys.argv[1]
+    path = sys.argv[2]
+    editor = os.environ.get("EDITOR", "$EDITOR")
+    repl = CustomRepl(ns, path, editor)
+    repl.setup()
+    repl.show_help()
+
+
+def main() -> None:
+    """Entrypoint, should be replaced by a `Omni.Cli.main`."""
+    if sys.argv[1] == "test":
+        test()
+    else:
+        move()
+
+
+if __name__ == "__main__":
+    main()