pythonstartup

"""
Python Startup File

Author:           Brandon Perez <bmperez@alumni.cmu.edu>
Creation Date:    Tuesday, February 16, 2016 at 11:19:26 AM EDT

Description:

    This file is executed whenever an interactive Python session starts up and sets up settings for the REPL.

    This file sets up persistent command history, tab completion, and default imported modules for the session.
"""


# ----------------------------------------------------------------------------------------------------------------------
# Default Modules
# ----------------------------------------------------------------------------------------------------------------------

# Import the Automatic Imports module so it's available in the interpreter, if it is installed. This module lazily
# imports all available Python modules. The modules are available, but the actual importing doesn't happen until the
# module is used in the interpreter. This lazy importing is also recursive. This must come before any other imports.
try:
    from autoimp import *   # type: ignore[import]  # noqa: F401, F403
except ImportError:
    print('Unable to import "autoimp" module. Run `python -m pip install autoimp` to install it.')

# Import the See module if it has been installed on this system.
try:
    import see as _see  # type: ignore[import]
except ImportError:
    print('Unable to import "see" module. Run `python -m pip install see` to install it.')

# Import the various container types and iteration tools without their qualified nammes, import the Path class so it can
# be easily used, and import the typing module for type annotations.
from collections import *   # noqa: F401, F403
from itertools import *     # noqa: F401, F403
from pathlib import Path
import typing


# ----------------------------------------------------------------------------------------------------------------------
# Global Namespace Values
# ----------------------------------------------------------------------------------------------------------------------

def _add_pint_units() -> None:
    """Adds the common units from the Pint module to the global namespace for the REPL.

    This uses the pint module, instantiates the various units through its interface, and then adds them to the global
    namespace so the user can access them from the REPL. This adds units like pint, meter, and all their various
    aliases.
    """

    # General Imports
    try:
        import pint as units_pint
    except ImportError:
        print('Unable to import the "pint" module. Run `python -m pip install pint` to install it.')
        return

    # The pint unit registry generates unit classes on the fly by overriding the `__getattr__` method and adds them as
    # class attributes as they are referenced. The complete list of units is stored under the `_unit`, attribute as a
    # dictionary, so this is used get the attributes and generate the classes. Not all of the units are in ASCII and
    # they will fail if they are referenced, so they are skipped. This also doesn't cover all units, some SI unit
    # prefixes like milliliters aren't listed in the unit list, so these have to be manually referenced.
    _ureg = units_pint.UnitRegistry()
    _ureg.milliliter
    globals().update({name: getattr(_ureg, name) for name in _ureg._units.keys() if name.isascii()})


def _add_module_aliases() -> None:
    """Adds some common aliases for modules to the global namespace so the user doesn't have to import them this way.

    This allows the user in the REPL to use these aliases to refer to these modules automatically via the `autoimp`
    module, such as `np` for `numpy`.
    """

    global_namespace = globals()
    if 'tensorflow' in global_namespace:
        global_namespace['tf'] = global_namespace['tensorflow']
    if 'numpy' in globals():
        global_namespace['np'] = global_namespace['numpy']


_add_pint_units()
_add_module_aliases()
del _add_pint_units, _add_module_aliases

# ----------------------------------------------------------------------------------------------------------------------
# Readline Settings
# ----------------------------------------------------------------------------------------------------------------------

# The history entry that was selected by the reverse history search, to be displayed on the subsequent line.
_SELECTED_HISTORY_ENTRY: str = ''


def _setup_readline() -> None:
    """Sets up tab completion, persistent command history, and shortcuts for the interactive session.

    This file imports the settings from a user's inputrc file and also sets up a shortcut to leverage FZF to enable
    fuzzy searching over the command history.
    """

    # The Readline library is only supported on Linux and Mac systems.
    import platform
    if platform.system().lower() != 'linux':
        return

    # Standard Imports
    import atexit
    import readline

    # Enable tab completion to be used in the Python interpreter.
    readline.parse_and_bind('tab: complete')

    # Set the persistent command history to be unlimited in size.
    readline.set_history_length(-1)

    # Read the contents of the Inputrc initialization file to get the standard Readline settings.
    inputrc_path = Path.home() / '.inputrc'
    readline.read_init_file(inputrc_path)

    # Read the contents of the Python history file if one currently exists to get the command history.
    history_path = Path.home() / '.python_history'
    if history_path.exists():
        readline.read_history_file(history_path)

    # When the interpreter exits, write the new history to the history file to enable persistent command history. Note
    # that this also includes the history that was read from the file, so this is effectively an append operation.
    atexit.register(readline.write_history_file, history_path)

    # Replace the standard reverse history search with a version that uses the FZF program. This works by invoking a
    # Python function to do the search, then using another function for the pre-input hook to place the selected entry
    # on the next prompt line.
    readline.parse_and_bind(r'"\C-r": "\C-e\C-u_fzf_history_search()\C-m"')
    readline.set_pre_input_hook(_readline_pre_input_hook)


def _fzf_history_search() -> None:
    """Perform a search over the Python command history using the FZF program (a fuzzy finder).

    This function is triggered when Ctrl-R is typed by the user. After the search is complete, this function updates a
    global variable that is utilized by the pre-input hook function to place the selected history on the next prompt
    line. This is required because there is no way to edit the current line buffer natively using Python.
    """

    # General Imports
    import os
    import subprocess

    # Readline Imports
    import readline

    # Declare the global variables accessed by this function.
    global _SELECTED_HISTORY_ENTRY

    # Remove this function call from the Python command history so that it doesn't pollute the history.
    history_index = readline.get_current_history_length() - 1
    readline.remove_history_item(history_index)

    # Read the entire history into a string.
    history_generator = (readline.get_history_item(i + 1) for i in range(readline.get_current_history_length()))
    history_entries = '\n'.join(history_generator)

    # Source the FZF CTRL-R and default options and combine them with standard options to form the FZF options.
    fzf_ctrl_r_opts = os.getenv('FZF_CTRL_R_OPTS', default='')
    fzf_default_opts = os.getenv('FZF_DEFAULT_OPTS', default='')
    fzf_options = (f'--height 40% {fzf_default_opts} --tac --sync --tiebreak=index --bind=ctrl-r:toggle-sort'
                   f' {fzf_ctrl_r_opts} --multi')

    # Run FZF to perform a fuzzy search over the history entries, passing the options in the environment variable.
    fzf_env = os.environ.copy()
    fzf_env['FZF_DEFAULT_OPTS'] = fzf_options
    process = subprocess.run(['fzf'], input = history_entries, stdout = subprocess.PIPE, env = fzf_env,
                             encoding = 'utf-8')

    # Remove the final newline from the input returned by the program, so the cursor is on the final line.
    _SELECTED_HISTORY_ENTRY = process.stdout.rstrip('\n')


def _readline_pre_input_hook() -> None:
    """The pre-input hook for readline, which is called after the prompt is printed, but before user input is entered.

    This function simply inserts the selected history entry from the reverse history search, if one was performed on the
    previous line into the readline buffer. This is equivalent to if the user typed it themselves.
    """

    # Readline Imports
    import readline

    # Declare the global variables accessed by this function.
    global _SELECTED_HISTORY_ENTRY

    # Insert the selected history entry into the current line. If a search wasn't performed, the string is empty.
    readline.insert_text(_SELECTED_HISTORY_ENTRY)
    readline.redisplay()

    # Unconditionally reset the selected history entry to an empty string
    _SELECTED_HISTORY_ENTRY = ''


# Remove setup Readline function so it doesn't pollute the interactive session's namespace.
_setup_readline()
del _setup_readline


# ----------------------------------------------------------------------------------------------------------------------
# Prompt Settings
# ----------------------------------------------------------------------------------------------------------------------

def _setup_prompt() -> None:
    """Modifies the PS1 and PS2 prompts used for the interactive session.

    If possible, terminal color codes are used. The PS1 prompt shows the Python version while the PS2 prompt aligns line
    continuations with the same level as the PS1 prompt offsets the first line.
    """

    # General Imports
    import curses
    import sys

    # Generate the color strings to use in the format strings.
    normal = _ansi_escape_sequence('sgr0')
    green = _ansi_escape_sequence('setaf', color = curses.COLOR_GREEN)
    yellow = _ansi_escape_sequence('setaf', color = curses.COLOR_YELLOW)
    magenta_background = _ansi_escape_sequence('setab', color = curses.COLOR_MAGENTA)

    # Override the default prompt used in the interactive session, printed at the beginning of each line, and set it up
    # to display the version of Python being used.
    version = sys.version_info
    sys.ps1 = f'{green}[{yellow}Python-{version.major}.{version.minor}.{version.micro}{green}]>>> {normal}'

    # Override the default prompt used for line continuations, so that multi-line commands can be copied properly. Align
    # the lines with the PS1 prompt length, so function definitions look nice.
    ps1_len = len(f'[Python-{version.major}.{version.minor}.{version.micro}]>>> ')
    ps2_len = ps1_len - 1
    space = ' '
    sys.ps2 = f'{magenta_background}{space:<{ps2_len}}{normal} '


def _ansi_escape_sequence(capability: str, color: typing.Optional[int] = None) -> str:
    """Gets the ANSI escape sequence that represents the given capability for the terminal.

    This function will get the escape sequence if the terminal supports the capability or modifier (and optional color
    code if it is a foreground or background coloring one) and then convert it to a UTF-8 string.

    Args:
        capability:     The capability to get the code for, which is typically the text modifier for the terminal.
        color:          The color code to use for the capability, only necessary if it is a foreground or background
                        setting capability.

    Returns:
        The ANSI escape sequence as a string representing the capability and optional color for the terminal or an empty
        string if the terminal does not support that capability.
    """

    # If Python is not being run on a Linux system, then the curses library is unavailable, so colors cannot be used.
    import platform
    if platform.system().lower() != 'linux':
        return ''

    # Setup the terminal for use by the curses library.
    import curses
    curses.setupterm()

    # Check that the terminal supports the requested capability and get it if it does.
    cap_code = curses.tigetstr(capability)
    if cap_code is None:
        return ''

    # If a color code was not specified, then the capability string is used as is. Otherwise, we need to get the full
    # color code string.
    if color is None:
        terminal_code = cap_code.decode('utf-8')
    else:
        terminal_code = curses.tparm(cap_code, color).decode('utf-8')

    # Note that the \001 and \002 delimiters tell Readline not to count the characters towards the line length. This is
    # necessary to wrap long lines properly.
    return f'\001{terminal_code}\002'


# Remove the setup prompt function so it doesn't pollute the interactive session's namespace.
_setup_prompt()
del _setup_prompt


# ----------------------------------------------------------------------------------------------------------------------
# Functions for the REPL
# ----------------------------------------------------------------------------------------------------------------------

def see(obj: object, pattern: typing.Optional[str] = None) -> _see.SeeResult:
    """Gets the attributes of the given object and optionally filters it using a glob or regex pattern.

    This serves as a better replacement for the `dir` function in Python and this overrides the default see.see function
    to allow for a pattern to be optionally specified. This is a bit more convenient than doing
    `see(obj).filter(pattern)`.
    convenient than doing see(obj).filter(pattern). Note that all members of a object begin with a '.' character, so
    this needs to be in the pattern. Regular expressions can be specified if the string begins with a forward-slash '/'.

    Args:
        obj:        The object to see the attributes for.
        pattern:    The pattern to pass to the filter method. This can be a glob or a regular expression. Note that
                    standard attributes of an object begin with '.' and regular expressions must start with a
                    forward-slash.

    Returns:
        A SeeResult object which represents the list of attributes for the object, optionally filtered with the pattern.

    Raises:
        RuntimeError:   The see module was not imported and thus isn't installed.
    """

    # If the See module could not be imported, then this function cannot run.
    if '_see' not in globals():
        raise RuntimeError('The see module is not installed so the see function cannot be used.')

    # Based on if a pattern is specified or not, invoke the appropriate functions.
    if (pattern is None):
        return _see.see(obj)
    else:
        return _see.see(obj).filter(pattern)


def wildcard_import_file(python_file_path: Path) -> None:
    """Emulates a wild card import of the Python file at the given path, if it exists.

    This follows the behavior of a wild card import (i.e. `from module import *`). If the module defines `__all__`, then
    only those names are imported. Otherwise, all names are imported which do not start with an underscore.

    Args:
        python_file_path:   The path to the Python file to import.
    """

    # General Imports
    import importlib
    import importlib.util

    # Otherwise, dynamically import the module with import lib's facilities.
    python_file_name = python_file_path.name.replace(python_file_path.suffix, '')
    module_loader = importlib.machinery.SourceFileLoader(python_file_name, str(python_file_path))
    module_spec = importlib.util.spec_from_loader(module_loader.name, module_loader)

    assert(module_spec is not None)
    assert(module_spec.loader is not None)
    module = importlib.util.module_from_spec(module_spec)
    module_spec.loader.exec_module(module)

    # Emulate a wild card import by importing all user-defined or specified variables into the global namespace. This is
    # either the `__all__` attribute, or all names that don't start with an underscore if `__all__` is not defined.
    default_import_variables = [name for name in module.__dict__.keys() if not name.startswith('_')]
    module_user_defined = getattr(module, '__all__', default_import_variables)
    module_variables = {name: getattr(module, name) for name in module_user_defined}
    globals().update(module_variables)


# ----------------------------------------------------------------------------------------------------------------------
# Machine-Local Settings
# ----------------------------------------------------------------------------------------------------------------------

def _import_local_startup() -> None:
    """Imports the local Python startup file, if one exists."""

    local_pythonstartup_path = Path.home() / '.pythonstartup_local'
    if local_pythonstartup_path.exists():
        wildcard_import_file(local_pythonstartup_path)  # noqa: NEWS100


# Import the local Python startup file if one exists.
_import_local_startup()
del _import_local_startup