[ipython] [refactor] Misc tweaks to make #1308 easier to review (#1584)

* [Misc] Use 'dill.source' of 'inspect' to run codes in interactive shell

* [skip ci] Add dill to CI

* _ShellInspectorWrapper

* Hack IDLE to make it happy

* Fix IDLE

* [skip ci] cache: we do care user experience!!

* improve stability

* [skip ci] fix example exit

* [skip ci] fix exit in IPython

* [skip ci] improve comment

* fix exec risk in ti debug (@rexwangcc)

* [skip ci] Fix ti debug too verbose

* Better line

* fix test_cli

* [skip ci] idle_hacker.py

* fix typo and improve hacker

* [skip ci] reimprov

* improve idle inspector

* [skip ci] add tag [IPython]

* revert idle

* [skip ci] Revert "revert idle"

This reverts commit 794bf05.

* tmp save

* well

* clean

* [skip ci] re

* [skip ci] al

* t

* improve hack

* nice message

* [skip ci] enforce code format

* ti idle_hacker -r

* [skip ci] enforce code format

* fix ipython

* revert

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>

.gitignore

@@ -57,6 +57,7 @@ __pycache__
 *.ini
 *.egg-info
 .tlang_cache
+.tidle_*
 /taichi/common/version.h
 /taichi/common/commit_hash.h
 /python/test_env

misc/idle_hello.py

@@ -0,0 +1,9 @@
+import taichi as ti
+
+
+@ti.kernel
+def func():
+    pass
+
+
+func()

python/taichi/lang/shell.py

@@ -1,4 +1,4 @@
-import sys, os
+import sys, os, atexit
 
 
 class ShellType:
@@ -9,95 +9,111 @@ class ShellType:
     SCRIPT = None
 
 
-def get_shell_name():
-    """
-    Detect which type of shell is using.
-    Can be IPython, IDLE, Python native, or none.
-    """
-    shell = os.environ.get('TI_SHELL_TYPE')
-    if shell is not None:
-        return getattr(ShellType, shell.upper())
-
-    try:
-        import __main__ as main
-        if hasattr(main, '__file__'):  # Called from a script?
-            return ShellType.SCRIPT
-    except:
-        pass
-
-    # Let's detect which type of interactive shell is being used.
-    # As you can see, huge engineering efforts are done here just to
-    # make IDLE and IPython happy. Hope our users really love them :)
-
-    try:  # IPython / Jupyter?
-        return 'IPython ' + get_ipython().__class__.__name__
-    except:
-        # Note that we can't simply do `'IPython' in sys.modules`,
-        # since it seems `torch` will import IPython on it's own too..
-        if hasattr(__builtins__, '__IPYTHON__'):
-            return ShellType.IPYBASED
-
-    try:
-        if getattr(sys, 'ps1', sys.flags.interactive):
-            return ShellType.NATIVE
-    except:
-        pass
-
-    return ShellType.SCRIPT
-
-
 class ShellInspectorWrapper:
     """
     Wrapper of the `inspect` module. When interactive shell detected,
     we will redirect getsource() calls to the corresponding inspector
     provided by / suitable for each type of shell.
     """
-    def __init__(self):
-        self.name = get_shell_name()
-
-        if self.name is not None:
-            print('[Taichi] Interactive shell detected:', self.name)
-
-        if self.name is None:
+    @staticmethod
+    def get_shell_name(exclude_script=False):
+        """
+        Detect which type of shell is using.
+        Can be IPython, IDLE, Python native, or none.
+        """
+        shell = os.environ.get('TI_SHELL_TYPE')
+        if shell is not None:
+            return getattr(ShellType, shell.upper())
+
+        if not exclude_script:
+            try:
+                import __main__ as main
+                if hasattr(main, '__file__'):  # Called from a script?
+                    return ShellType.SCRIPT
+            except:
+                pass
+
+        # Let's detect which type of interactive shell is being used.
+        # As you can see, huge engineering efforts are done here just to
+        # make IDLE and IPython happy. Hope our users really love them :)
+
+        try:  # IPython / Jupyter?
+            return 'IPython ' + get_ipython().__class__.__name__
+        except:
+            # Note that we can't simply do `'IPython' in sys.modules`,
+            # since it seems `torch` will import IPython on it's own too..
+            if hasattr(__builtins__, '__IPYTHON__'):
+                return ShellType.IPYBASED
+
+        try:
+            if getattr(sys, 'ps1', sys.flags.interactive):
+                return ShellType.NATIVE
+        except:
+            pass
+
+        return ShellType.SCRIPT
+
+    @staticmethod
+    def create_inspector(name):
+        if name is None:
             # `inspect` for "Python script"
             import inspect
-            self.getsource = inspect.getsource
-            self.getsourcelines = inspect.getsourcelines
-            self.getsourcefile = inspect.getsourcefile
+            return inspect
 
-        elif self.name == ShellType.NATIVE:
+        elif name == ShellType.NATIVE:
             # `dill.source` for "Python native shell"
             try:
                 import dill
             except ImportError as e:
                 raise ImportError(
                     'In order to run Taichi in Python interactive shell, '
                     'Please execute `python3 -m pip install --user dill`')
-            self.getsource = dill.source.getsource
-            self.getsourcelines = dill.source.getsourcelines
-            self.getsourcefile = dill.source.getsourcefile
+            return dill.source
 
-        elif self.name.startswith('IPython'):
+        elif name.startswith('IPython'):
             # `IPython.core.oinspect` for "IPython advanced shell"
-            def getsource(o):
-                import IPython
-                return IPython.core.oinspect.getsource(o)
+            return IPythonInspectorWrapper()
 
-            def getsourcelines(o):
-                import IPython
-                lineno = IPython.core.oinspect.find_source_lines(o)
-                lines = IPython.core.oinspect.getsource(o).split('\n')
-                return lines, lineno
+        else:
+            raise RuntimeError(f'Shell type "{name}" not supported')
 
-            def getsourcefile(o):
-                return '<IPython>'
+    def __init__(self):
+        self.name = self.get_shell_name()
+        if self.name is not None:
+            print(f'[Taichi] Interactive shell detected: {self.name}')
 
-            self.getsource = getsource
-            self.getsourcelines = getsourcelines
-            self.getsourcefile = getsourcefile
+        self.inspector = self.create_inspector(self.name)
 
-        else:
-            raise RuntimeError(f'Shell type "{self.name}" not supported')
+    def getsource(self, o):
+        return self.inspector.getsource(o)
+
+    def getsourcelines(self, o):
+        return self.inspector.getsourcelines(o)
+
+    def getsourcefile(self, o):
+        return self.inspector.getsourcefile(o)
+
+
+class IPythonInspectorWrapper:
+    """`inspect` module wrapper for IPython / Jupyter notebook"""
+    def __init__(self):
+        pass
+
+    def getsource(self, o):
+        import IPython
+        return IPython.core.oinspect.getsource(o)
+
+    def getsourcelines(self, o):
+        import IPython
+        lineno = IPython.core.oinspect.find_source_lines(o)
+        lines = IPython.core.oinspect.getsource(o).split('\n')
+        return lines, lineno
+
+    def getsourcefile(self, o):
+        import IPython
+        lineno = IPython.core.oinspect.find_source_lines(o)
+        return f'<IPython:{lineno}>'
 
 
 oinspect = ShellInspectorWrapper()
+# TODO: also detect print according to shell type

python/taichi/main.py

@@ -966,7 +966,7 @@ def debug(self, arguments: list = sys.argv[2:]):
         ti.core.set_core_trigger_gdb_when_crash(True)
         os.environ['TI_DEBUG'] = '1'
 
-        runpy.run_path(args.filename)
+        runpy.run_path(args.filename, run_name='__main__')
 
     @register
     def task(self, arguments: list = sys.argv[2:]):