Revert "revert"

This reverts commit 202c235.

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())

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'<IPython:{lineno}>'
 
 
+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 '<IDLE>'
+
+
 oinspect = ShellInspectorWrapper()
 # TODO: also detect print according to shell type

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"""