Add a wrapper for inspect in JIT to produce better error message (#25415)

Dmytro Dzhulgakov · facebook-github-bot · commit df338f80a631 · 2019-09-14T21:27:51.000-07:00
Summary: If source code is not available due to packaging (e.g. sources are compiled to .pyc), TorchScript produces very obscure error message. This tries to make it nicer and allow to customize message by overriding _utils_internal. Pull Request resolved: #25415 Test Plan: Really hard to unittest properly. Did one off testing by compiling to .pyc and checking the message. Differential Revision: D17118238 Pulled By: dzhulgakov fbshipit-source-id: 3cbfee0abddc8613000680548bfe0b8ed52a36b0
diff --git a/torch/_jit_internal.py b/torch/_jit_internal.py
@@ -9,6 +9,7 @@
 import warnings
 import torch._C
 from torch._six import builtins
+from torch._utils_internal import get_source_lines_and_file
 
 # Wrapper functions that can call either of 2 functions depending on a boolean
 # argument
@@ -477,11 +478,11 @@ def _get_overloaded_methods(method, mod_class):
     if overloads is None:
         return None
 
-    method_line_no = inspect.getsourcelines(method)[1]
-    mod_class_fileno = inspect.getsourcelines(mod_class)[1]
-    mod_end_fileno = mod_class_fileno + len(inspect.getsourcelines(mod_class)[0])
+    method_line_no = get_source_lines_and_file(method)[1]
+    mod_class_fileno = get_source_lines_and_file(mod_class)[1]
+    mod_end_fileno = mod_class_fileno + len(get_source_lines_and_file(mod_class)[0])
     if not (method_line_no >= mod_class_fileno and method_line_no <= mod_end_fileno):
-        raise Exception("Overloads are not useable when a module is redaclared within the same file: " + str(method))
+        raise Exception("Overloads are not useable when a module is redeclared within the same file: " + str(method))
     return overloads
 
 try:
diff --git a/torch/_utils_internal.py b/torch/_utils_internal.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import os
+import inspect
 
 # this arbitrary-looking assortment of functionality is provided here
 # to have a central place for overrideable behavior. The motivating
@@ -33,5 +34,23 @@ def resolve_library_path(path):
     return os.path.realpath(path)
 
 
+def get_source_lines_and_file(obj):
+    """
+    Wrapper around inspect.getsourcelines and inspect.getsourcefile.
+
+    Returns: (sourcelines, file_lino, filename)
+    """
+    filename = None  # in case getsourcefile throws
+    try:
+        filename = inspect.getsourcefile(obj)
+        sourcelines, file_lineno = inspect.getsourcelines(obj)
+    except OSError as e:
+        raise OSError((
+            "Can't get source for {}. TorchScript requires source access in order to carry out compilation. " +
+            "Make sure original .py files are available. Original error: {}").format(filename, e))
+
+    return sourcelines, file_lineno, filename
+
+
 TEST_MASTER_ADDR = '127.0.0.1'
 TEST_MASTER_PORT = 29500
diff --git a/torch/jit/annotations.py b/torch/jit/annotations.py
@@ -9,6 +9,7 @@
 from torch._C import TensorType, TupleType, FloatType, IntType, \
     ListType, StringType, DictType, BoolType, OptionalType, ClassType
 from textwrap import dedent
+from torch._utils_internal import get_source_lines_and_file
 
 
 PY35 = sys.version_info >= (3, 5)
@@ -46,7 +47,7 @@ def get_signature(fn):
 
     type_line, source = None, None
     try:
-        source = dedent(inspect.getsource(fn))
+        source = dedent(''.join(get_source_lines_and_file(fn)[0]))
         type_line = get_type_line(source)
     except TypeError:
         pass
@@ -63,7 +64,7 @@ def get_signature(fn):
 # a function takes.
 def get_num_params(fn, loc):
     try:
-        source = dedent(inspect.getsource(fn))
+        source = dedent(''.join(get_source_lines_and_file(fn)[0]))
     except (TypeError, IOError):
         return None
     if source is None:
diff --git a/torch/jit/frontend.py b/torch/jit/frontend.py
@@ -7,6 +7,7 @@
 from textwrap import dedent
 from torch._six import PY2
 from torch._C._jit_tree_views import *
+from torch._utils_internal import get_source_lines_and_file
 
 # Borrowed from cPython implementation
 # https://github.com/python/cpython/blob/561612d8456cfab5672c9b445521113b847bd6b3/Lib/textwrap.py#L411#
@@ -146,9 +147,8 @@ def get_jit_class_def(cls, self_name):
     method_defs = [get_jit_def(method[1],
                    self_name=self_name) for method in methods]
 
-    sourcelines, file_lineno = inspect.getsourcelines(cls)
+    sourcelines, file_lineno, filename = get_source_lines_and_file(cls)
     source = ''.join(sourcelines)
-    filename = inspect.getsourcefile(cls)
     dedent_src = dedent(source)
     py_ast = ast.parse(dedent_src)
     leading_whitespace_len = len(source.split('\n', 1)[0]) - len(dedent_src.split('\n', 1)[0])
@@ -157,9 +157,8 @@ def get_jit_class_def(cls, self_name):
 
 
 def get_jit_def(fn, self_name=None):
-    sourcelines, file_lineno = inspect.getsourcelines(fn)
+    sourcelines, file_lineno, filename = get_source_lines_and_file(fn)
     source = ''.join(sourcelines)
-    filename = inspect.getsourcefile(fn)
     dedent_src = dedent(source)
     py_ast = ast.parse(dedent_src)
     if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef):
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -1,5 +1,4 @@
 import difflib
-import inspect
 import os
 import io
 import shutil
@@ -12,6 +11,7 @@
 from contextlib import closing, contextmanager
 from ._utils import _import_dotted_name
 from ._six import string_classes as _string_classes
+from torch._utils_internal import get_source_lines_and_file
 if sys.version_info[0] == 2:
     import cPickle as pickle
 else:
@@ -285,8 +285,8 @@ def persistent_id(obj):
             serialized_container_types[obj] = True
             source_file = source = None
             try:
-                source_file = inspect.getsourcefile(obj)
-                source = inspect.getsource(obj)
+                source_lines, _, source_file = get_source_lines_and_file(obj)
+                source = ''.join(obj)
             except Exception:  # saving the source is optional, so we can ignore any errors
                 warnings.warn("Couldn't retrieve source code for container of "
                               "type " + obj.__name__ + ". It won't be checked "
@@ -449,7 +449,7 @@ def restore_location(storage, location):
 
     def _check_container_source(container_type, source_file, original_source):
         try:
-            current_source = inspect.getsource(container_type)
+            current_source = ''.join(get_source_lines_and_file(container_type)[0])
         except Exception:  # saving the source is optional, so we can ignore any errors
             warnings.warn("Couldn't retrieve source code for container of "
                           "type " + container_type.__name__ + ". It won't be checked "
