Merge pull request #1250 from google/google_sync

Google sync

pytype/load_pytd.py

@@ -5,7 +5,7 @@
 import os
 import pickle
 
-from typing import Dict, Iterable, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from pytype import module_utils
 from pytype import utils
@@ -49,6 +49,7 @@ class ResolvedModule:
   module_name: str
   filename: str
   ast: pytd.TypeDeclUnit
+  metadata: List[str]
 
 
 class Module:
@@ -66,18 +67,20 @@ class Module:
     pickle: The AST as a pickled string. As long as this field is not None, the
       ast will be None.
     has_unresolved_pointers: Whether all ClassType pointers have been filled in
+    metadata: The metadata extracted from the picked file.
   """
 
   _INIT_NAMES = ("__init__.pyi", f"__init__.{pytd_utils.PICKLE_EXT}")
 
   # pylint: disable=redefined-outer-name
-  def __init__(self, module_name, filename, ast,
-               pickle=None, has_unresolved_pointers=True):
+  def __init__(self, module_name, filename, ast, metadata=None, pickle=None,
+               has_unresolved_pointers=True):
     self.module_name = module_name
     self.filename = filename
     self.ast = ast
     self.pickle = pickle
     self.has_unresolved_pointers = has_unresolved_pointers
+    self.metadata = metadata or []
   # pylint: enable=redefined-outer-name
 
   def needs_unpickling(self):
@@ -160,7 +163,7 @@ def get_resolved_modules(self) -> Dict[str, ResolvedModule]:
     for name, mod in self._modules.items():
       if not mod.has_unresolved_pointers:
         resolved_modules[name] = ResolvedModule(
-            mod.module_name, mod.filename, mod.ast)
+            mod.module_name, mod.filename, mod.ast, mod.metadata)
     return resolved_modules
 
   def _base_modules(self):
@@ -800,7 +803,8 @@ def load_file(self, module_name, filename, mod_ast=None):
     dependencies = {d: names for d, names in loaded_ast.dependencies
                     if d != loaded_ast.ast.name}
     loaded_ast = serialize_ast.EnsureAstName(loaded_ast, module_name, fix=True)
-    self._modules[module_name] = Module(module_name, filename, loaded_ast.ast)
+    self._modules[module_name] = Module(
+        module_name, filename, loaded_ast.ast, metadata=loaded_ast.metadata)
     self._load_ast_dependencies(dependencies, lookup_ast=mod_ast,
                                 lookup_ast_name=module_name)
     try:

pytype/tools/xref/indexer_test.py

@@ -25,13 +25,15 @@ def index_code(self, code, **kwargs):
       options.tweak(**args)
       return indexer.process_file(options, preserve_pytype_vm=True)
 
-  def generate_kythe(self, code):
+  def generate_kythe(self, code, **kwargs):
     """Generate a kythe index from a code string."""
     with file_utils.Tempdir() as d:
       d.create_file("t.py", code)
       options = config.Options.create(d["t.py"])
       options.tweak(pythonpath=[d.path], version=self.python_version)
-      kythe_args = kythe.Args(corpus="corpus", root="root", path="path")
+      k_args = {k: k for k in ["corpus", "root", "path"]}
+      k_args.update(kwargs)
+      kythe_args = kythe.Args(**k_args)
       ix = indexer.process_file(options)
       kg = kythe.generate_graph(ix, kythe_args)
       # Collect all the references from the kythe graph.
@@ -196,6 +198,19 @@ def f(x):
     node = kythe_index[3]
     self.assertEqual(node["source"]["language"], "python")
 
+  def test_kythe_skip_stdlib(self):
+    code = textwrap.dedent("""
+      import os
+    """)
+    kythe_index = self.generate_kythe(code)
+    imp = [x for x in kythe_index
+           if x.get("edge_kind") == "/kythe/edge/ref/imports"]
+    self.assertEqual(len(imp), 1)
+    kythe_index = self.generate_kythe(code, skip_stdlib=True)
+    imp = [x for x in kythe_index
+           if x.get("edge_kind") == "/kythe/edge/ref/imports"]
+    self.assertFalse(imp)
+
   def test_multiline_attr(self):
     # Test that lookahead doesn't crash.
     self.index_code(textwrap.dedent("""

pytype/tools/xref/kythe.py

@@ -15,6 +15,7 @@ class Args:
   root: str
   corpus: str
   path: str
+  skip_stdlib: bool = False
 
 
 # Kythe nodes
@@ -52,10 +53,12 @@ def __init__(self, source, args=None):
       self.root = args.root
       self.corpus = args.corpus
       self.path = args.path or source.filename
+      self.skip_stdlib = args.skip_stdlib
     else:
       self.root = ""
       self.corpus = ""
       self.path = source.filename
+      self.skip_stdlib = False
     self.entries = []
     self._seen_entries = set()
     self.file_vname = self._add_file(source.text)
@@ -90,12 +93,12 @@ def _add_entry(self, entry):
     self._seen_entries.add(entry)
     self.entries.append(entry)
 
-  def vname(self, signature, filepath=None):
+  def vname(self, signature, filepath=None, root=None):
     return VName(
         signature=signature,
         path=filepath or self.path,
         language="python",
-        root=self.root,
+        root=root or self.root,
         corpus=self.corpus)
 
   def stdlib_vname(self, signature, filepath=None):
@@ -166,8 +169,9 @@ def _process_deflocs(kythe: Kythe, index: indexer.Indexer):
         pass
       else:
         alias_vname = _make_defn_vname(kythe, index, alias)
-        kythe.add_edge(
-            source=defn_vname, target=alias_vname, edge_name="aliases")
+        if alias_vname:
+          kythe.add_edge(
+              source=defn_vname, target=alias_vname, edge_name="aliases")
 
       # Emit a docstring if we have one.
       doc = defn.doc
@@ -202,6 +206,7 @@ def _make_defn_vname(kythe, index, defn):
   if isinstance(defn, indexer.Remote):
     remote = defn.module
     if remote in index.resolved_modules:
+      is_generated = "generated" in index.resolved_modules[remote].metadata
       if remote in index.imports:
         # The canonical path from the imports_map is the best bet for
         # module->filepath translation.
@@ -216,8 +221,14 @@ def _make_defn_vname(kythe, index, defn):
       else:
         sig = "module." + defn.name
       if path.startswith("pytd:"):
+        if kythe.skip_stdlib:
+          # Skip builtin and stdlib imports since we don't have a filepath.
+          # TODO(mdemello): Link to the typeshed definition
+          return None
         return kythe.stdlib_vname(
             sig, "pytd:" + index.resolved_modules[remote].module_name)
+      elif is_generated:
+        return kythe.vname(sig, path, root="root/genfiles")
       else:
         return kythe.vname(sig, path)
     else:

pytype/tools/xref/parse_args.py

@@ -42,6 +42,10 @@ def add_kythe_field(parser, field):
       "--imports_info", type=str, action="store",
       dest="imports_info", default=None,
       help="Information for mapping import .pyi to files. ")
+  # Don't index builtins and stdlib.
+  parser.add_argument("--skip-stdlib", action="store_true",
+                      dest="skip_stdlib", default=None,
+                      help="Display inferred types.")
   # Add options from pytype-single.
   wrapper = datatypes.ParserWrapper(parser)
   pytype_config.add_basic_options(wrapper)
@@ -76,5 +80,7 @@ def parse_args(argv):
   pytype_options = pytype_config.Options(cli_args, command_line=True)
   pytype_options.tweak(**parser.get_pytype_kwargs(args))
   kythe_args = kythe.Args(
-      corpus=args.kythe_corpus, root=args.kythe_root, path=args.kythe_path)
+      corpus=args.kythe_corpus, root=args.kythe_root, path=args.kythe_path,
+      skip_stdlib=args.skip_stdlib
+  )
   return (args, kythe_args, pytype_options)