diff --git a/.gitignore b/.gitignore index da03592f06274..34cb8517bcab3 100644 --- a/.gitignore +++ b/.gitignore @@ -57,6 +57,7 @@ __pycache__ *.ini *.egg-info .tlang_cache +.tidle_* /taichi/common/version.h /taichi/common/commit_hash.h /python/test_env diff --git a/misc/idle_hello.py b/misc/idle_hello.py new file mode 100644 index 0000000000000..8e49844690ce2 --- /dev/null +++ b/misc/idle_hello.py @@ -0,0 +1,9 @@ +import taichi as ti + + +@ti.kernel +def func(): + pass + + +func() diff --git a/python/taichi/lang/shell.py b/python/taichi/lang/shell.py index 9e0d3ea94f7e7..9f84241888801 100644 --- a/python/taichi/lang/shell.py +++ b/python/taichi/lang/shell.py @@ -1,4 +1,4 @@ -import sys, os +import sys, os, atexit class ShellType: @@ -9,63 +9,58 @@ 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 @@ -73,31 +68,52 @@ def __init__(self): 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 '' + 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'' oinspect = ShellInspectorWrapper() +# TODO: also detect print according to shell type diff --git a/python/taichi/main.py b/python/taichi/main.py index c4e06a4001659..6e633b0561157 100644 --- a/python/taichi/main.py +++ b/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:]):