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