From 3b76f609a1016d784029f28de505eed85aa61e8b Mon Sep 17 00:00:00 2001 From: mdemello Date: Tue, 28 Jun 2022 16:36:06 -0700 Subject: [PATCH] BEGIN_PUBLIC Add some metadata fields to the serialized ast. Added fields: src_path: Full path to the original source file. metadata: List of arbitrary string-encoded metadata END_PUBLIC The new fields are mainly for kythe, so that cross references to generated code link to the correct place. Also includes a pytype release, so that kythe stays in sync. PiperOrigin-RevId: 457845537 --- pytype/pytd/serialize_ast.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/pytype/pytd/serialize_ast.py b/pytype/pytd/serialize_ast.py index d8fb858c7..c18ef6354 100644 --- a/pytype/pytd/serialize_ast.py +++ b/pytype/pytd/serialize_ast.py @@ -63,7 +63,9 @@ class SerializableTupleClass(NamedTuple): dependencies: List[Tuple[str, Set[str]]] late_dependencies: List[Tuple[str, Set[str]]] class_type_nodes: Optional[List[pytd.ClassType]] - is_package: bool + is_package: bool # TODO(rechen): remove this, can be computed from src_path + src_path: Optional[str] + metadata: List[str] class SerializableAst(SerializableTupleClass): @@ -83,11 +85,15 @@ class SerializableAst(SerializableTupleClass): visited and have their .cls set. If this attribute is None the whole AST will be visited and all found ClassType instances will have their .cls set. + is_package: True if the original source file was a package __init__. + src_path: Optionally, the filepath of the original source file. + metadata: A list of arbitrary string-encoded metadata. """ Replace = SerializableTupleClass._replace # pylint: disable=no-member,invalid-name -def StoreAst(ast, filename=None, open_function=open, is_package=False): +def StoreAst(ast, filename=None, open_function=open, is_package=False, + src_path=None, metadata=None): """Loads and stores an ast to disk. Args: @@ -96,6 +102,8 @@ def StoreAst(ast, filename=None, open_function=open, is_package=False): function instead returns the pickled string. open_function: A custom file opening function. is_package: Whether the module with the given ast is a package. + src_path: Optionally, the filepath of the original source file. + metadata: A list of arbitrary string-encoded metadata. Returns: The pickled string, if no filename was given. (None otherwise.) @@ -116,10 +124,15 @@ def StoreAst(ast, filename=None, open_function=open, is_package=False): indexer = FindClassTypesVisitor() ast.Visit(indexer) ast = ast.Visit(visitors.CanonicalOrderingVisitor()) + + metadata = metadata or [] + return pytd_utils.SavePickle( SerializableAst( ast, sorted(dependencies.items()), sorted(late_dependencies.items()), - sorted(indexer.class_type_nodes), is_package=is_package), + sorted(indexer.class_type_nodes), is_package=is_package, + src_path=src_path, metadata=metadata, + ), filename, open_function=open_function)