diff --git a/python/taichi/idle_hacker.py b/python/taichi/idle_hacker.py new file mode 100644 index 0000000000000..175775aa14834 --- /dev/null +++ b/python/taichi/idle_hacker.py @@ -0,0 +1,144 @@ +''' +A dirty hack injector for python/taichi/lang/shell.py:IDLEInspectorWrapper +''' + +import os +import functools +import taichi as ti + +our_code = "__import__('taichi.idle_hacker').idle_hacker.hack(InteractiveInterpreter)" + + +def get_taichi_dir(): + return os.path.dirname(os.path.abspath(ti.__file__)) + + +def get_ipc_file(pid): + return os.path.join(get_taichi_dir(), '.tidle_' + str(pid)) + + +def get_backup_file(): + return os.path.join(get_taichi_dir(), '.tidle_backup.code.py') + + +def idle_ipc_write(source): + with open(get_ipc_file(os.getpid()), 'a') as f: + f.write('\n===\n' + source) + + +def hack(InteractiveInterpreter): + old_runsource = InteractiveInterpreter.runsource + + @functools.wraps(old_runsource) + def new_runsource(self, source, *args, **kwargs): + idle_ipc_write(source) + return old_runsource(self, source, *args, **kwargs) + + InteractiveInterpreter.runsource = new_runsource + + +def show_error(): + try: + import code + path = code.__file__ + except: + path = '/usr/lib/python3.8/code.py' + + print('''Hi! Dear Taichi user: + + It's detected that you are using Python IDLE in **interactive mode**. + However, Taichi could not be fully functional due to IDLE limitation, sorry :( + Either run Taichi in IDLE file mode, or use IPython / Jupyter instead. + But we do care about your experience, no matter which shell you prefer to use. + So, in order to play Taichi with your favorite IDLE, we may do a dirty hack: + Open "{path}" and append the following line to the buttom of this file: + +''' + f' {our_code}' + ''' + +If you don't find where to append, we offer a script to inject the code: +''') + + if ti.get_os_name() == 'win': + print(' python3 -m taichi idle_hacker') + else: + print(' sudo python3 -m taichi idle_hacker') + + print(''' + Then, restart IDLE and enjoy, the sky is blue and we are wizards! +''') + + +def startup_clean(): + filename = get_ipc_file(os.getppid()) + try: + os.unlink(filename) + except: + pass + else: + import taichi as ti + ti.info(f'File "{filename}" cleaned') + + +def read_ipc_file(): + # The IDLE GUI and Taichi is running in separate process, + # So we have to create temporary files for portable IPC :( + filename = get_ipc_file(os.getppid()) + try: + with open(filename) as f: + src = f.read() + except FileNotFoundError as e: + show_error() + src = '' + return src + + +def main(revert=False): + import code + import shutil + + print('Injection dest:', code.__file__) + + backup_file = get_backup_file() + + if not revert: + with open(code.__file__) as f: + if our_code in f.read(): + print('ERROR: Taichi hack code already exists') + return 1 + + if not os.path.exists(backup_file): + shutil.copy(code.__file__, backup_file) + print(f'Backup saved in {backup_file}') + + print('Appending our hack code...') + + with open(code.__file__, 'a') as f: + f.write('\n' + our_code) + + print('Done, thank for trusting!') + + else: + with open(code.__file__) as f: + if our_code not in f.read(): + print('ERROR: Taichi hack code not exist') + return 1 + + if not os.path.exists(backup_file): + print(f'ERROR: Backup file {backup_file} not found!') + print( + 'Sorry, please consider manually remove the code or reinstall IDLE :(' + ) + return 1 + + print('Moving backup file to dest...') + shutil.move(backup_file, code.__file__) + + print('Done, sorry for the trouble!') + + return 0 + + +if __name__ == '__main__': + exit(main()) diff --git a/python/taichi/lang/shell.py b/python/taichi/lang/shell.py index 9f84241888801..07d6e7edc5a59 100644 --- a/python/taichi/lang/shell.py +++ b/python/taichi/lang/shell.py @@ -6,6 +6,7 @@ class ShellType: IPYTHON = 'IPython TerminalInteractiveShell' JUPYTER = 'IPython ZMQInteractiveShell' IPYBASED = 'IPython Based Shell' + IDLE = 'Python IDLE shell' SCRIPT = None @@ -45,6 +46,9 @@ def get_shell_name(exclude_script=False): if hasattr(__builtins__, '__IPYTHON__'): return ShellType.IPYBASED + if 'idlelib' in sys.modules: + return ShellType.IDLE + try: if getattr(sys, 'ps1', sys.flags.interactive): return ShellType.NATIVE @@ -74,6 +78,10 @@ def create_inspector(name): # `IPython.core.oinspect` for "IPython advanced shell" return IPythonInspectorWrapper() + elif name == ShellType.IDLE: + # `.tidle_xxx` for "Python IDLE shell" + return IDLEInspectorWrapper() + else: raise RuntimeError(f'Shell type "{name}" not supported') @@ -84,12 +92,46 @@ def __init__(self): self.inspector = self.create_inspector(self.name) + if hasattr(self.inspector, 'startup_clean'): + self.inspector.startup_clean() + + def try_reset_shell_type(self): + new_name = self.get_shell_name(exclude_script=True) + if self.name != new_name: + print( + f'[Taichi] Shell type changed from "{self.name}" to "{new_name}"' + ) + + self.name = new_name + self.inspector = self.create_inspector(self.name) + + def _catch_forward(foo): + """ + If Taichi starts within IDLE file mode, and after that user moved to interactive mode, + then there will be an OSError, since it's switched from None to Python IDLE shell... + We have to reset the shell type and create a corresponding inspector at this moment. + """ + import functools + + @functools.wraps(foo) + def wrapped(self, *args, **kwargs): + try: + return foo(self, *args, **kwargs) + except OSError: + self.try_reset_shell_type() + return foo(self, *args, **kwargs) + + return wrapped + + @_catch_forward def getsource(self, o): return self.inspector.getsource(o) + @_catch_forward def getsourcelines(self, o): return self.inspector.getsourcelines(o) + @_catch_forward def getsourcefile(self, o): return self.inspector.getsourcefile(o) @@ -115,5 +157,62 @@ def getsourcefile(self, o): return f'' +class IDLEInspectorWrapper: + """`inspect` module wrapper for IDLE / Blender scripting module""" + + # Thanks to IDLE's lack of support with `inspect`, + # we have to use a dirty hack to support Taichi there. + + def __init__(self): + self.idle_cache = {} + from taichi.idle_hacker import startup_clean + self.startup_clean = startup_clean + + def getsource(self, o): + func_id = id(o) + if func_id in self.idle_cache: + return self.idle_cache[func_id] + + from taichi.idle_hacker import read_ipc_file + src = read_ipc_file() + + # If user added our 'hacker-code' correctly, + # then the content of `.tidle_xxx` should be: + # + # === + # import taichi as ti + # + # === + # @ti.kernel + # def func(): # x.find('def ') locate to here + # pass + # + # === + # func() + # + + func_name = o.__name__ + for x in reversed(src.split('===')): + x = x.strip() + i = x.find('def ') + if i == -1: + continue + name = x[i + 4:].split(':', maxsplit=1)[0] + name = name.split('(', maxsplit=1)[0] + if name.strip() == func_name: + self.idle_cache[func_name] = x + return x + else: + raise NameError(f'Could not find source for {func_name}!') + + def getsourcelines(self, o): + lineno = 2 # TODO: consider include lineno in .tidle_xxx? + lines = self.getsource(o).split('\n') + return lines, lineno + + def getsourcefile(self, o): + return '' + + oinspect = ShellInspectorWrapper() # TODO: also detect print according to shell type diff --git a/python/taichi/main.py b/python/taichi/main.py index 6e633b0561157..ee94ea57be141 100644 --- a/python/taichi/main.py +++ b/python/taichi/main.py @@ -968,6 +968,22 @@ def debug(self, arguments: list = sys.argv[2:]): runpy.run_path(args.filename, run_name='__main__') + @register + def idle_hacker(self, arguments: list = sys.argv[2:]): + """Run idle hack code injector""" + parser = argparse.ArgumentParser( + prog='ti idle_hacker', description=f"{self.idle_hacker.__doc__}") + parser.add_argument( + '-r', + '--revert', + dest='revert', + action='store_true', + help='Revert the hack injection in case you meet troubles') + args = parser.parse_args(arguments) + + from .idle_hacker import main + return main(revert=args.revert) + @register def task(self, arguments: list = sys.argv[2:]): """Run a specific task"""