fix pydruid.console prompt_toolkit + pygments integration

[erhlee-bird]

With both the latest pydruid pip release and the version installed from the repo HEAD, the CLI does not function properly. Tested with Python 3.8.7 and 3.9.1.

latest pip release

$ pip uninstall pydruid
$ pip install pydruid[cli]
$ pydruid

Traceback (most recent call last):
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/cache.py", line 38, in get
    return self._data[key]
KeyError: ('', 94043350247808)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ubuntu/.pyenv/versions/3.8.7/bin/pydruid", line 33, in <module>
    sys.exit(load_entry_point('pydruid==0.6.2', 'console_scripts', 'pydruid')())
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/pydruid/console.py", line 171, in main
    query = prompt(
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/shortcuts/prompt.py", line 1384, in prompt
    return session.prompt(
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/shortcuts/prompt.py", line 1013, in prompt
    return self.app.run(set_exception_handler=set_exception_handler)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 825, in run
    return loop.run_until_complete(
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/asyncio/base_events.py", line 616, in run_until_complete
    return future.result()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 792, in run_async
    return await _run_async2()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 774, in _run_async2
    result = await _run_async()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 710, in _run_async
    self._redraw()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 543, in _redraw
    self.context.run(run_in_context)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 529, in run_in_context
    self.renderer.render(self, self.layout)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/renderer.py", line 637, in render
    layout.container.preferred_height(size.columns, size.rows).preferred,
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 325, in preferred_height
    dimensions = [
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 326, in <listcomp>
    c.preferred_height(width, max_available_height) for c in self._all_children
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 795, in preferred_height
    return self.content.preferred_height(width, max_available_height)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 325, in preferred_height
    dimensions = [
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 326, in <listcomp>
    c.preferred_height(width, max_available_height) for c in self._all_children
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 2598, in preferred_height
    return self.content.preferred_height(width, max_available_height)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 1623, in preferred_height
    return self._merge_dimensions(
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 1651, in _merge_dimensions
    preferred = get_preferred()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 1616, in preferred_content_height
    return self.content.preferred_height(
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 647, in preferred_height
    content = self.create_content(width, height=1)  # Pass a dummy '1' as height.
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 779, in create_content
    get_processed_line = self._create_get_processed_line_func(
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 730, in _create_get_processed_line_func
    return create_func()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 717, in create_func
    get_line = self._get_formatted_text_for_line_func(document)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 678, in _get_formatted_text_for_line_func
    return self._fragment_cache.get(key, get_formatted_text_for_line)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/cache.py", line 41, in get
    value = getter_func()
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 675, in get_formatted_text_for_line
    return self.lexer.lex_document(document)
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/lexers/base.py", line 78, in lex_document
    return lexer.lex_document(document)
AttributeError: type object 'SqlLexer' has no attribute 'lex_document'

HEAD

The same error appears installing from HEAD.

$ pip uninstall pydruid
$ pip install git+https://github.com/druid-io/pydruid
$ pydruid

Traceback (most recent call last):                                                                                                                                          
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/cache.py", line 38, in get                                                            
    return self._data[key]                                                                                                                                                  
KeyError: ('', 94297281620784)                                                                                                                                              
                                                                                                                                                                            
During handling of the above exception, another exception occurred:                                                                                                         
                                                                                                                                                                            
Traceback (most recent call last):                                                                                                                                          
  File "/home/ubuntu/.pyenv/versions/3.8.7/bin/pydruid", line 33, in <module>                                                                                               
    sys.exit(load_entry_point('pydruid==0.999.0.dev0', 'console_scripts', 'pydruid')())                                                                                     
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/pydruid/console.py", line 171, in main                                                               
    query = prompt(                                                                                                                                                         
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/shortcuts/prompt.py", line 1384, in prompt                                            
    return session.prompt(                                                                                                                                                  
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/shortcuts/prompt.py", line 1013, in prompt                                            
    return self.app.run(set_exception_handler=set_exception_handler)                                                                                                        
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 825, in run                                         
    return loop.run_until_complete(                                                                                                                                         
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/asyncio/base_events.py", line 616, in run_until_complete                                                           
    return future.result()                                                                                                                                                  
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 792, in run_async                                   
    return await _run_async2()                                                                                                                                              
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 774, in _run_async2                                 
    result = await _run_async()                                                                                                                                             
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 710, in _run_async                                  
    self._redraw()                                                                                                                                                          
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 543, in _redraw                                     
    self.context.run(run_in_context)                                                                                                                                        
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/application/application.py", line 529, in run_in_context                              
    self.renderer.render(self, self.layout)                                                                                                                                 
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/renderer.py", line 637, in render                                                     
    layout.container.preferred_height(size.columns, size.rows).preferred,                                                                                                   
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 325, in preferred_height                                  
    dimensions = [                                                                                                                                                          
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 326, in <listcomp>                                        
    c.preferred_height(width, max_available_height) for c in self._all_children                                                                                             
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 795, in preferred_height                                  
    return self.content.preferred_height(width, max_available_height)                                                                                                       
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 325, in preferred_height                                  
    dimensions = [                                                                                                                                                          
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 326, in <listcomp>                                        
    c.preferred_height(width, max_available_height) for c in self._all_children                                                                                             
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 2598, in preferred_height                                 
    return self.content.preferred_height(width, max_available_height)                                                                                                       
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 1623, in preferred_height                                 
    return self._merge_dimensions(                                                                                                                                          
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 1651, in _merge_dimensions                                
    preferred = get_preferred()                                                                                                                                             
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/containers.py", line 1616, in preferred_content_height                         
    return self.content.preferred_height(                                                                                                                                   
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 647, in preferred_height                                    
    content = self.create_content(width, height=1)  # Pass a dummy '1' as height.                                                                                           
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 779, in create_content                                      
    get_processed_line = self._create_get_processed_line_func(                                                                                                              
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 730, in _create_get_processed_line_func                     
    return create_func()                                                                                                                                                    
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 717, in create_func                                         
    get_line = self._get_formatted_text_for_line_func(document)                                                                                                             
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 678, in _get_formatted_text_for_line_func                   
    return self._fragment_cache.get(key, get_formatted_text_for_line)                                                                                                       
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/cache.py", line 41, in get                                                            
    value = getter_func()                                                                                                                                                   
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/layout/controls.py", line 675, in get_formatted_text_for_line                         
    return self.lexer.lex_document(document)                                                                                                                                
  File "/home/ubuntu/.pyenv/versions/3.8.7/lib/python3.8/site-packages/prompt_toolkit/lexers/base.py", line 78, in lex_document                                             
    return lexer.lex_document(document)                                                                                                                                     
AttributeError: type object 'SqlLexer' has no attribute 'lex_document'                                                

This PR

$ pip uninstall pydruid
$ pip install git+https://github.com/erhlee-bird/pydruid
$ pydruid

>                                                                                                                                                                            
GoodBye!

Fix

The lexer and style classes are pygments-based and just needed some pygments-specific wrappers to function properly.

See:

• https://python-prompt-toolkit.readthedocs.io/en/master/pages/tutorials/repl.html
• https://python-prompt-toolkit.readthedocs.io/en/master/pages/advanced_topics/styling.html#loading-a-style-from-pygments