From ea0c0c4af44b66e9de22f5e6e6df31f4e8e60a93 Mon Sep 17 00:00:00 2001 From: Zhao Liang Date: Wed, 11 Jan 2023 19:37:40 +0800 Subject: [PATCH] [Lang] Fix cannot use taichi in REPL (#7114) This PR fixes the issue: #7111 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- python/taichi/lang/_wrap_inspect.py | 106 +++++++++++++++++++--------- setup.py | 3 +- 2 files changed, 74 insertions(+), 35 deletions(-) diff --git a/python/taichi/lang/_wrap_inspect.py b/python/taichi/lang/_wrap_inspect.py index a8d420c6ba455..77bd66e92a994 100644 --- a/python/taichi/lang/_wrap_inspect.py +++ b/python/taichi/lang/_wrap_inspect.py @@ -1,13 +1,33 @@ +# Taichi's custom inspect module. +# This module is used by Taichi's ast transformer to parse the source code. +# Currently this module is aimed for working in the following modes: +# 1. Usual Python/IPython mode, e.g. python script.py +# In this case we mainly rely on the built-in `inspect` module, except +# we need some hacks when we are in IPython mode and there is a cell magic. +# 2. Blender's scripting mode, e.g. Users write Taichi code in the scripting +# window in Blender and press the run button. In this case we need to +# retrieve the source using Blender's `bpy.data.texts` and write it to a temp +# file so that the inspect module can parse. +# 3. The interactive shell mode, e.g. Users directly type their code in the +# interactive shell. In this case we use `dill` to get the source. +# +# NB: Running Taichi in other modes are likely not supported. + import atexit import inspect import os import tempfile +import dill + _builtin_getfile = inspect.getfile _builtin_findsource = inspect.findsource def _find_source_with_custom_getfile_func(func, obj): + """Use a custom function `func` to replace inspect's `getfile`, return the + source found by the new routine and restore the original `getfile` back. + """ inspect.getfile = func # replace with our custom func source = inspect.findsource(obj) inspect.getfile = _builtin_getfile # restore @@ -15,12 +35,19 @@ def _find_source_with_custom_getfile_func(func, obj): def _blender_get_text_name(filename: str): + """Extract filename from path in the Blender mode.""" + # In Blender's scripting mode, unsaved files are named + # like `/Text`, `/Text.001`, `/test.py`, etc. + # We simply remove this path seperator. if filename.startswith(os.path.sep) and filename.count(os.path.sep) == 1: return filename[1:] # "/Text.001" --> "Text.001" + # Saved text files are named like `some-path/xxx.blend/Text` or + # `some-path/xxx.blend/test.py` + # We drop the path and extract the filename with extension. index = filename.rfind('.blend' + os.path.sep) if index != -1: - return filename[index + 7:] # "hello.blend/test.py" --> "test.py" + return filename[index + 7:] # "xxx.blend/test.py" --> "test.py" return None @@ -28,21 +55,26 @@ def _blender_get_text_name(filename: str): def _blender_findsource(obj): try: import bpy # pylint: disable=import-outside-toplevel - except ImportError: - raise IOError('Not in Blender environment!') + except: + raise ImportError('Not in Blender environment!') + # Inspect's built-in `getfile` returns the filename like + # `/Text`, `/Text.001`, `some-path/xxx.blend/test.py` + # This filename may not be a full valid path. filename = _builtin_getfile(obj) + # Extract the text name without path text_name = _blender_get_text_name(filename) if text_name is None: raise IOError( 'Object `{obj.__name__}` is not defined in a .blend file!') - + # Get the lines of code via text_name lines = bpy.data.texts[text_name].as_string() - # Now we have found the filename and code lines. + # Now we have found the lines of code. # We first check if they are already cached, to avoid file io in each query. try: filename = _blender_findsource._saved_inspect_cache[lines] # pylint: disable=no-member except KeyError: + # Save the code to a valid path. fd, filename = tempfile.mkstemp(prefix='_Blender_', suffix=f'_{text_name}.py') os.close(fd) @@ -53,6 +85,7 @@ def _blender_findsource(obj): _blender_findsource._saved_inspect_cache[lines] = filename # pylint: disable=no-member atexit.register(os.unlink, filename) # Remove file when program exits + # Our custom getfile function def wrapped_getfile(ob): if id(ob) == id(obj): return filename @@ -67,24 +100,29 @@ def wrapped_getfile(ob): def _Python_IPython_findsource(obj): try: - # In Python and IPython the builtin findsource would suffice in most cases + # In Python and IPython the builtin inspect would suffice in most cases return _builtin_findsource(obj) except IOError: # Except that the cell has a magic command like %%time or %%timeit - # In this case the filename returned by getfile is wrong + # In this case the filename returned by the built-in's getfile is wrong, + # it becomes something like `` or ``. filename = _builtin_getfile(obj) if (filename in {"", ""}): try: ip = get_ipython() if ip is not None: + # So we are in IPython's cell magic session_id = ip.history_manager.get_last_session_id() fd, filename = tempfile.mkstemp(prefix='_IPython_', suffix=f'_{session_id}.py') os.close(fd) - # The latest lines of code are stored in this file + # The latest lines of code can be retrived from here lines = ip.history_manager._i00 - # Remove the magic command (and spaces/sep around it) before saving to a file + # `lines` is a string that also contains the cell magic + # command, we need to remove the magic command + # (and spaces/sep around it) to obtain a valid Python code + # snippet before saving it to a file index = lines.find("%time") lines_stripped = lines[index:] lines_stripped = lines_stripped.split(maxsplit=1)[1] @@ -98,9 +136,16 @@ def _Python_IPython_findsource(obj): func = lambda obj: filename return _find_source_with_custom_getfile_func(func, obj) - except: + except ImportError: pass - raise IOError(f"Cannot find source code for Object: {obj}") + raise IOError( + f"Cannot find source code for Object: {obj}, it's likely \ +you are not running Taichi from command line or IPython.") + + +def _REPL_findsource(obj): + """Findsource in the interactive shell mode.""" + return dill.source.findsource(obj) def _custom_findsource(obj): @@ -108,9 +153,16 @@ def _custom_findsource(obj): return _Python_IPython_findsource(obj) except IOError: try: - return _blender_findsource(obj) + return _REPL_findsource(obj) except: - raise IOError(f"Cannot find source code for Object: {obj} ") + try: + return _blender_findsource(obj) + except: + raise IOError( + f"Cannot find source code for Object: {obj}, this \ +is possibly because of you are running Taichi in an environment that Taichi's own \ +inspect module cannot find the source. Please report an issue to help us fix: \ +https://github.com/taichi-dev/taichi/issues") class _InspectContextManager: @@ -123,30 +175,16 @@ def __exit__(self, *_): def getsourcelines(obj): - try: - with _InspectContextManager(): - return inspect.getsourcelines(obj) - except: - raise IOError(f"Cannot get the source lines of {obj}. \ - This is possibly because of you are running Taichi in an environment \ - in which Taichi's own inspect module cannot find the source file. \ - Please report an issue to help us fix this problem: https://github.com/taichi-dev/taichi/issues" - ) + with _InspectContextManager(): + return inspect.getsourcelines(obj) def getsourcefile(obj): - try: - with _InspectContextManager(): - ret = inspect.getsourcefile(obj) - if ret is None: - ret = inspect.getfile(obj) - return ret - except: - raise IOError(f"Cannot get the source file of {obj}. \ - This is possibly because of you are running Taichi in an environment \ - in which Taichi's own inspect module cannot find the source file. \ - Please report an issue to help us fix this problem: https://github.com/taichi-dev/taichi/issues" - ) + with _InspectContextManager(): + ret = inspect.getsourcefile(obj) + if ret is None: + ret = inspect.getfile(obj) + return ret __all__ = ['getsourcelines', 'getsourcefile'] diff --git a/setup.py b/setup.py index 75cb33e3d2db2..2e9dc5493d667 100644 --- a/setup.py +++ b/setup.py @@ -177,7 +177,8 @@ def cmake_install_manifest_filter(manifest_files): url='https://github.com/taichi-dev/taichi', python_requires=">=3.6,<3.11", install_requires=[ - 'numpy', 'colorama', 'rich', 'astunparse;python_version<"3.9"' + 'numpy', 'colorama', 'dill', 'rich', + 'astunparse;python_version<"3.9"' ], data_files=[(os.path.join('_lib', 'runtime'), data_files)], keywords=['graphics', 'simulation'],