diff options
Diffstat (limited to 'ptpython/lexer.py')
-rw-r--r-- | ptpython/lexer.py | 28 |
1 files changed, 28 insertions, 0 deletions
diff --git a/ptpython/lexer.py b/ptpython/lexer.py new file mode 100644 index 0000000..62e470f --- /dev/null +++ b/ptpython/lexer.py @@ -0,0 +1,28 @@ +from typing import Callable, Optional + +from prompt_toolkit.document import Document +from prompt_toolkit.formatted_text import StyleAndTextTuples +from prompt_toolkit.lexers import Lexer, PygmentsLexer +from pygments.lexers import BashLexer +from pygments.lexers import Python3Lexer as PythonLexer + +__all__ = ["PtpythonLexer"] + + +class PtpythonLexer(Lexer): + """ + Lexer for ptpython input. + + If the input starts with an exclamation mark, use a Bash lexer, otherwise, + use a Python 3 lexer. + """ + + def __init__(self, python_lexer: Optional[Lexer] = None) -> None: + self.python_lexer = python_lexer or PygmentsLexer(PythonLexer) + self.system_lexer = PygmentsLexer(BashLexer) + + def lex_document(self, document: Document) -> Callable[[int], StyleAndTextTuples]: + if document.text.startswith("!"): + return self.system_lexer.lex_document(document) + + return self.python_lexer.lex_document(document) |