path: root/Omni/Repl.py

#!/usr/bin/env run.sh
"""
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()