1from typing import Callable, Optional
2
3from prompt_toolkit.document import Document
4from prompt_toolkit.formatted_text import StyleAndTextTuples
5from prompt_toolkit.lexers import Lexer, PygmentsLexer
6from pygments.lexers import BashLexer
7from pygments.lexers import Python3Lexer as PythonLexer
8
9__all__ = ["PtpythonLexer"]
10
11
12class PtpythonLexer(Lexer):
13    """
14    Lexer for ptpython input.
15
16    If the input starts with an exclamation mark, use a Bash lexer, otherwise,
17    use a Python 3 lexer.
18    """
19
20    def __init__(self, python_lexer: Optional[Lexer] = None) -> None:
21        self.python_lexer = python_lexer or PygmentsLexer(PythonLexer)
22        self.system_lexer = PygmentsLexer(BashLexer)
23
24    def lex_document(self, document: Document) -> Callable[[int], StyleAndTextTuples]:
25        if document.text.startswith("!"):
26            return self.system_lexer.lex_document(document)
27
28        return self.python_lexer.lex_document(document)
29