projects.mako.cc - twitter-api-cdsw-solutions/blob - win_unicode_console/readline_hook.py

   1
   2 import sys, traceback
   3 from ctypes import pythonapi, cdll, c_size_t, c_char_p, c_void_p, cast, CFUNCTYPE, POINTER, addressof
   4
   5 PyMem_Malloc = pythonapi.PyMem_Malloc
   6 PyMem_Malloc.restype = c_size_t
   7 PyMem_Malloc.argtypes = [c_size_t]
   8
   9 strncpy = cdll.msvcrt.strncpy
  10 strncpy.restype = c_char_p
  11 strncpy.argtypes = [c_char_p, c_char_p, c_size_t]
  12
  13 HOOKFUNC = CFUNCTYPE(c_char_p, c_void_p, c_void_p, c_char_p)
  14
  15 PyOS_ReadlineFunctionPointer = c_void_p.in_dll(pythonapi, "PyOS_ReadlineFunctionPointer")
  16
  17
  18 def new_zero_terminated_string(b):
  19         p = PyMem_Malloc(len(b) + 1)
  20         strncpy(cast(p, c_char_p), b, len(b) + 1)
  21         return p
  22
  23
  24 class ReadlineHookManager:
  25         def __init__(self):
  26                 self.readline_wrapper_ref = HOOKFUNC(self.readline_wrapper)
  27                 self.address = c_void_p.from_address(addressof(self.readline_wrapper_ref)).value
  28                 self.original_address = PyOS_ReadlineFunctionPointer.value
  29                 self.readline_hook = None
  30
  31         def readline_wrapper(self, stdin, stdout, prompt):
  32                 try:
  33                         try:
  34                                 if sys.stdin.encoding != sys.stdout.encoding:
  35                                         raise ValueError("sys.stdin.encoding != sys.stdout.encoding, readline hook doesn't know, which one to use to decode prompt")
  36
  37                         except ValueError:
  38                                 traceback.print_exc(file=sys.stderr)
  39                                 try:
  40                                         prompt = prompt.decode("utf-8")
  41                                 except UnicodeDecodeError:
  42                                         prompt = ""
  43
  44                         else:
  45                                 prompt = prompt.decode(sys.stdout.encoding)
  46
  47                         try:
  48                                 line = self.readline_hook(prompt)
  49                         except KeyboardInterrupt:
  50                                 return 0
  51                         else:
  52                                 return new_zero_terminated_string(line.encode(sys.stdin.encoding))
  53
  54                 except:
  55                         print("Intenal win_unicode_console error", file=sys.stderr)
  56                         traceback.print_exc(file=sys.stderr)
  57                         return new_zero_terminated_string(b"\n")
  58
  59         def install_hook(self, hook):
  60                 self.readline_hook = hook
  61                 PyOS_ReadlineFunctionPointer.value = self.address
  62
  63         def restore_original(self):
  64                 self.readline_hook = None
  65                 PyOS_ReadlineFunctionPointer.value = self.original_address
  66
  67
  68 def readline(prompt):
  69         sys.stdout.write(prompt)
  70         sys.stdout.flush()
  71         return sys.stdin.readline()
  72
  73
  74 class PyReadlineManager:
  75         def __init__(self):
  76                 self.original_codepage = pyreadline.unicode_helper.pyreadline_codepage
  77
  78         def set_codepage(self, codepage):
  79                 pyreadline.unicode_helper.pyreadline_codepage = codepage
  80
  81         def restore_original(self):
  82                 self.set_codepage(self.original_codepage)
  83
  84 try:
  85         import pyreadline.unicode_helper
  86 except ImportError:
  87         pyreadline = None
  88 else:
  89         pyreadline_manager = PyReadlineManager()
  90
  91 manager = ReadlineHookManager()
  92
  93
  94 def enable(*, use_pyreadline=True):
  95         if use_pyreadline and pyreadline:
  96                 pyreadline_manager.set_codepage(sys.stdin.encoding)
  97                         # pyreadline assumes that encoding of all sys.stdio objects is the same
  98         else:
  99                 manager.install_hook(readline)
 100
 101 def disable():
 102         if pyreadline:
 103                 pyreadline_manager.restore_original()
 104
 105         manager.restore_original()
 106