[IPython] Source inspection dispatcher for better IDLE compatibility

docs/dev_install.rst

@@ -19,7 +19,7 @@ Installing Dependencies
 
   .. code-block:: bash
 
-    python3 -m pip install --user setuptools astpretty astor pybind11 Pillow
+    python3 -m pip install --user setuptools astpretty astor pybind11 Pillow dill
     python3 -m pip install --user pytest pytest-rerunfailures pytest-xdist yapf
     python3 -m pip install --user numpy GitPython coverage colorama autograd
 

examples/mpm128.py

@@ -105,9 +105,9 @@ def reset():
 gravity[None] = [0, -1]
 
 for frame in range(20000):
-  while gui.get_event(ti.GUI.PRESS):
+  if gui.get_event(ti.GUI.PRESS):
     if gui.event.key == 'r': reset()
-    elif gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: exit(0)
+    elif gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: break
   if gui.event is not None: gravity[None] = [0, 0] # if had any event
   if gui.is_pressed(ti.GUI.LEFT,  'a'): gravity[None][0] = -1
   if gui.is_pressed(ti.GUI.RIGHT, 'd'): gravity[None][0] = 1

examples/stable_fluid.py

@@ -244,10 +244,10 @@ def main():
     md_gen = MouseDataGen()
     paused = False
     while True:
-        while gui.get_event(ti.GUI.PRESS):
+        if gui.get_event(ti.GUI.PRESS):
             e = gui.event
             if e.key == ti.GUI.ESCAPE:
-                exit(0)
+                break
             elif e.key == 'r':
                 paused = False
                 reset()

misc/ci_setup.py

@@ -187,6 +187,7 @@ def run(self):
             "distro",
             "autograd",
             "astor",
+            "dill",
             "pytest",
             "pytest-xdist",
             "pytest-rerunfailures",

misc/prtags.json

@@ -27,5 +27,6 @@
   "mac"             : "Mac OS X",
   "windows"         : "Windows",
   "perf"            : "Performance improvements",
+  "ipython"         : "IPython and other shells",
   "release"         : "Release"
 }

python/taichi/lang/ast_checker.py

@@ -1,4 +1,5 @@
 import ast
+from .shell import oinspect
 
 
 class KernelSimplicityASTChecker(ast.NodeVisitor):
@@ -32,8 +33,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     def __init__(self, func):
         super().__init__()
         import inspect
-        self._func_file = inspect.getsourcefile(func)
-        self._func_lineno = inspect.getsourcelines(func)[1]
+        self._func_file = oinspect.getsourcefile(func)
+        self._func_lineno = oinspect.getsourcelines(func)[1]
         self._func_name = func.__name__
         self._scope_guards = []
 

python/taichi/lang/kernel.py

@@ -4,6 +4,7 @@
 import ast
 from .kernel_arguments import *
 from .util import *
+from .shell import oinspect
 import functools
 
 
@@ -55,7 +56,7 @@ def __call__(self, *args):
 
     def do_compile(self):
         from .impl import get_runtime
-        src = remove_indent(inspect.getsource(self.func))
+        src = remove_indent(oinspect.getsource(self.func))
         tree = ast.parse(src)
 
         func_body = tree.body[0]
@@ -75,7 +76,7 @@ def do_compile(self):
             print('After preprocessing:')
             print(astor.to_source(tree.body[0], indent_with='  '))
 
-        ast.increment_lineno(tree, inspect.getsourcelines(self.func)[1] - 1)
+        ast.increment_lineno(tree, oinspect.getsourcelines(self.func)[1] - 1)
 
         local_vars = {}
         #frame = inspect.currentframe().f_back
@@ -84,7 +85,7 @@ def do_compile(self):
         global_vars = copy.copy(self.func.__globals__)
         exec(
             compile(tree,
-                    filename=inspect.getsourcefile(self.func),
+                    filename=oinspect.getsourcefile(self.func),
                     mode='exec'), global_vars, local_vars)
         self.compiled = local_vars[self.func.__name__]
 
@@ -260,7 +261,7 @@ def materialize(self, key=None, args=None, arg_features=None):
         import taichi as ti
         ti.trace("Compiling kernel {}...".format(kernel_name))
 
-        src = remove_indent(inspect.getsource(self.func))
+        src = remove_indent(oinspect.getsource(self.func))
         tree = ast.parse(src)
         if self.runtime.print_preprocessed:
             import astor
@@ -300,7 +301,7 @@ def materialize(self, key=None, args=None, arg_features=None):
             print('After preprocessing:')
             print(astor.to_source(tree.body[0], indent_with='  '))
 
-        ast.increment_lineno(tree, inspect.getsourcelines(self.func)[1] - 1)
+        ast.increment_lineno(tree, oinspect.getsourcelines(self.func)[1] - 1)
 
         freevar_names = self.func.__code__.co_freevars
         closure = self.func.__closure__
@@ -316,7 +317,7 @@ def materialize(self, key=None, args=None, arg_features=None):
 
         exec(
             compile(tree,
-                    filename=inspect.getsourcefile(self.func),
+                    filename=oinspect.getsourcefile(self.func),
                     mode='exec'), global_vars, local_vars)
         compiled = local_vars[self.func.__name__]
 

python/taichi/lang/shell.py

@@ -0,0 +1,98 @@
+import sys, os, atexit
+
+
+class ShellType:
+    NATIVE = 'Python shell'
+    IPYTHON = 'IPython TerminalInteractiveShell'
+    JUPYTER = 'IPython ZMQInteractiveShell'
+    IPYBASED = 'IPython Based Shell'
+    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, wish our user 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:
+            # `inspect` for "Python script"
+            import inspect
+            self.getsource = inspect.getsource
+            self.getsourcelines = inspect.getsourcelines
+            self.getsourcefile = inspect.getsourcefile
+
+        elif self.name == ShellType.NATIVE:
+            # `dill.source` for "Python native shell"
+            import dill
+            self.getsource = dill.source.getsource
+            self.getsourcelines = dill.source.getsourcelines
+            self.getsourcefile = dill.source.getsourcefile
+
+        elif self.name.startswith('IPython'):
+            # `IPython.core.oinspect` for "IPython advanced shell"
+            def getsource(o):
+                import IPython
+                return IPython.core.oinspect.getsource(o)
+
+            def getsourcelines(o):
+                import IPython
+                lineno = IPython.core.oinspect.find_source_lines(o)
+                lines = IPython.core.oinspect.getsource(o).split('\n')
+                return lines, lineno
+
+            def getsourcefile(o):
+                return '<IPython>'
+
+            self.getsource = getsource
+            self.getsourcelines = getsourcelines
+            self.getsourcefile = getsourcefile
+
+        else:
+            raise RuntimeError(f'Shell type "{self.name}" not supported')
+
+
+oinspect = None

python/taichi/main.py

@@ -69,18 +69,18 @@ def __init__(self, debug: bool = False, test_mode: bool = False):
         parser.add_argument('command',
                             help="command from the above list to run")
 
-        # Print help if no command provided
-        if len(sys.argv[1:2]) == 0:
-            parser.print_help()
-            exit(1)
-
         # Flag for unit testing
         self.test_mode = test_mode
 
         self.main_parser = parser
 
     @timer
     def __call__(self):
+        # Print help if no command provided
+        if len(sys.argv[1:2]) == 0:
+            self.main_parser.print_help()
+            return 1
+
         # Parse the command
         args = self.main_parser.parse_args(sys.argv[1:2])
 
@@ -90,7 +90,7 @@ def __call__(self):
                 TaichiMain._exec_python_file(args.command)
             print(f"{args.command} is not a valid command!")
             self.main_parser.print_help()
-            exit(1)
+            return 1
 
         return getattr(self, args.command)(sys.argv[2:])
 
@@ -954,4 +954,4 @@ def main_debug():
 
 
 if __name__ == "__main__":
-    exit(main())
+    sys.exit(main())

setup.py

@@ -33,6 +33,7 @@
         'colorama',
         'setuptools',
         'astor',
+        'dill',
         # For testing:
         'pytest',
         'pytest-xdist',

tests/python/test_cli.py

@@ -22,20 +22,14 @@ def patch_sys_argv_helper(custom_argv: list):
 
 def test_cli_exit_one_with_no_command_provided():
     with patch_sys_argv_helper(["ti"]):
-        with pytest.raises(SystemExit) as pytest_wrapped_err:
-            cli = TaichiMain(test_mode=True)
-            cli()
-            assert pytest_wrapped_err.type == SystemExit
-            assert pytest_wrapped_err.value.code == 1
+        cli = TaichiMain(test_mode=True)
+        assert cli() == 1
 
 
 def test_cli_exit_one_with_bogus_command_provided():
     with patch_sys_argv_helper(["ti", "bogus-command-not-registered-yet"]):
-        with pytest.raises(SystemExit) as pytest_wrapped_err:
-            cli = TaichiMain(test_mode=True)
-            cli()
-            assert pytest_wrapped_err.type == SystemExit
-            assert pytest_wrapped_err.value.code == 1
+        cli = TaichiMain(test_mode=True)
+        assert cli() == 1
 
 
 def test_cli_can_dispatch_commands_to_methods_correctly():