Implement lazy loading mechanism for expensive metadata providers

libcst/_metadata_dependent.py

@@ -7,6 +7,7 @@
 from abc import ABC
 from contextlib import contextmanager
 from typing import (
+    Callable,
     cast,
     ClassVar,
     Collection,
@@ -31,6 +32,27 @@
 
 _UNDEFINED_DEFAULT = object()
 
+_SENTINEL = object()
+
+
+class LazyValue:
+    """
+    The class for implementing a lazy metadata loading mechanism that improves the
+    performance when retriving expensive metadata (e.g., qualified names). Providers
+    including :class:`~libcst.metadata.QualifiedNameProvider` use this class to load
+    the metadata of a certain node lazily when calling
+    :func:`~libcst.MetadataDependent.get_metadata`.
+    """
+
+    def __init__(self, callable: Callable[[], _T]) -> None:
+        self.callable = callable
+        self.return_value: object = _SENTINEL
+
+    def __call__(self) -> object:
+        if self.return_value is _SENTINEL:
+            self.return_value = self.callable()
+        return self.return_value
+
 
 class MetadataDependent(ABC):
     """
@@ -107,6 +129,9 @@ def get_metadata(
             )
 
         if default is not _UNDEFINED_DEFAULT:
-            return cast(_T, self.metadata[key].get(node, default))
+            value = self.metadata[key].get(node, default)
         else:
-            return cast(_T, self.metadata[key][node])
+            value = self.metadata[key][node]
+        if isinstance(value, LazyValue):
+            value = value()
+        return cast(_T, value)

libcst/codemod/visitors/_apply_type_annotations.py

@@ -5,7 +5,7 @@
 
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import cast, Collection, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import libcst as cst
 import libcst.matchers as m
@@ -17,7 +17,7 @@
 from libcst.codemod.visitors._gather_imports import GatherImportsVisitor
 from libcst.codemod.visitors._imports import ImportItem
 from libcst.helpers import get_full_name_for_node
-from libcst.metadata import PositionProvider, QualifiedNameProvider
+from libcst.metadata import PositionProvider, QualifiedName, QualifiedNameProvider
 
 
 NameOrAttribute = Union[cst.Name, cst.Attribute]
@@ -48,7 +48,12 @@ def _get_unique_qualified_name(
     visitor: m.MatcherDecoratableVisitor, node: cst.CSTNode
 ) -> str:
     name = None
-    names = [q.name for q in visitor.get_metadata(QualifiedNameProvider, node)]
+    names = [
+        q.name
+        for q in cast(
+            Collection[QualifiedName], visitor.get_metadata(QualifiedNameProvider, node)
+        )
+    ]
     if len(names) == 0:
         # we hit this branch if the stub is directly using a fully
         # qualified name, which is not technically valid python but is

libcst/codemod/visitors/_gather_string_annotation_names.py

@@ -9,7 +9,7 @@
 import libcst.matchers as m
 from libcst.codemod._context import CodemodContext
 from libcst.codemod._visitor import ContextAwareVisitor
-from libcst.metadata import MetadataWrapper, QualifiedNameProvider
+from libcst.metadata import MetadataWrapper, QualifiedName, QualifiedNameProvider
 
 FUNCS_CONSIDERED_AS_STRING_ANNOTATIONS = {"typing.TypeVar"}
 
@@ -45,7 +45,9 @@ def leave_Annotation(self, original_node: cst.Annotation) -> None:
         self._annotation_stack.pop()
 
     def visit_Call(self, node: cst.Call) -> bool:
-        qnames = self.get_metadata(QualifiedNameProvider, node)
+        qnames = cast(
+            Collection[QualifiedName], self.get_metadata(QualifiedNameProvider, node)
+        )
         if any(qn.name in self._typing_functions for qn in qnames):
             self._annotation_stack.append(node)
             return True

libcst/matchers/_matcher_base.py

@@ -31,6 +31,7 @@
 import libcst
 import libcst.metadata as meta
 from libcst import FlattenSentinel, MaybeSentinel, RemovalSentinel
+from libcst._metadata_dependent import LazyValue
 
 
 class DoNotCareSentinel(Enum):
@@ -1544,7 +1545,11 @@ def _fetch(provider: meta.ProviderT, node: libcst.CSTNode) -> object:
         if provider not in metadata:
             metadata[provider] = wrapper.resolve(provider)
 
-        return metadata.get(provider, {}).get(node, _METADATA_MISSING_SENTINEL)
+        node_metadata = metadata.get(provider, {}).get(node, _METADATA_MISSING_SENTINEL)
+        if isinstance(node_metadata, LazyValue):
+            node_metadata = node_metadata()
+
+        return node_metadata
 
     return _fetch
 

libcst/metadata/name_provider.py

@@ -5,10 +5,10 @@
 
 import dataclasses
 from pathlib import Path
-from typing import Collection, List, Mapping, Optional, Union
+from typing import cast, Collection, List, Mapping, Optional, Union
 
 import libcst as cst
-from libcst._metadata_dependent import MetadataDependent
+from libcst._metadata_dependent import LazyValue, MetadataDependent
 from libcst.helpers.module import calculate_module_and_package, ModuleNameAndPackage
 from libcst.metadata.base_provider import BatchableMetadataProvider
 from libcst.metadata.scope_provider import (
@@ -17,8 +17,10 @@
     ScopeProvider,
 )
 
+_UNDEFINED_DEFAULT = object
 
-class QualifiedNameProvider(BatchableMetadataProvider[Collection[QualifiedName]]):
+
+class QualifiedNameProvider(BatchableMetadataProvider[_UNDEFINED_DEFAULT]):
     """
     Compute possible qualified names of a variable CSTNode
     (extends `PEP-3155 <https://www.python.org/dev/peps/pep-3155/>`_).
@@ -64,7 +66,10 @@ def has_name(
         visitor: MetadataDependent, node: cst.CSTNode, name: Union[str, QualifiedName]
     ) -> bool:
         """Check if any of qualified name has the str name or :class:`~libcst.metadata.QualifiedName` name."""
-        qualified_names = visitor.get_metadata(QualifiedNameProvider, node, set())
+        qualified_names = cast(
+            Collection[QualifiedName],
+            visitor.get_metadata(QualifiedNameProvider, node, set()),
+        )
         if isinstance(name, str):
             return any(qn.name == name for qn in qualified_names)
         else:
@@ -78,7 +83,9 @@ def __init__(self, provider: "QualifiedNameProvider") -> None:
     def on_visit(self, node: cst.CSTNode) -> bool:
         scope = self.provider.get_metadata(ScopeProvider, node, None)
         if scope:
-            self.provider.set_metadata(node, scope.get_qualified_names_for(node))
+            self.provider.set_metadata(
+                node, LazyValue(lambda: scope.get_qualified_names_for(node))
+            )
         else:
             self.provider.set_metadata(node, set())
         super().on_visit(node)
@@ -171,7 +178,10 @@ def __init__(
         self.provider = provider
 
     def on_visit(self, node: cst.CSTNode) -> bool:
-        qnames = self.provider.get_metadata(QualifiedNameProvider, node)
+        qnames = cast(
+            Collection[QualifiedName],
+            self.provider.get_metadata(QualifiedNameProvider, node),
+        )
         if qnames is not None:
             self.provider.set_metadata(
                 node,

libcst/metadata/tests/test_name_provider.py

@@ -6,10 +6,11 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from textwrap import dedent
-from typing import Collection, Dict, Mapping, Optional, Set, Tuple
+from typing import cast, Collection, Dict, Mapping, Optional, Set, Tuple
 
 import libcst as cst
 from libcst import ensure_type
+from libcst._nodes.base import CSTNode
 from libcst.metadata import (
     FullyQualifiedNameProvider,
     MetadataWrapper,
@@ -22,11 +23,28 @@
 from libcst.testing.utils import data_provider, UnitTest
 
 
+class QNameVisitor(cst.CSTVisitor):
+
+    METADATA_DEPENDENCIES = (QualifiedNameProvider,)
+
+    def __init__(self) -> None:
+        self.qnames: Dict["CSTNode", Collection[QualifiedName]] = {}
+
+    def on_visit(self, node: cst.CSTNode) -> bool:
+        qname = cast(
+            Collection[QualifiedName], self.get_metadata(QualifiedNameProvider, node)
+        )
+        self.qnames[node] = qname
+        return True
+
+
 def get_qualified_name_metadata_provider(
     module_str: str,
 ) -> Tuple[cst.Module, Mapping[cst.CSTNode, Collection[QualifiedName]]]:
     wrapper = MetadataWrapper(cst.parse_module(dedent(module_str)))
-    return wrapper.module, wrapper.resolve(QualifiedNameProvider)
+    visitor = QNameVisitor()
+    wrapper.visit(visitor)
+    return wrapper.module, visitor.qnames
 
 
 def get_qualified_names(module_str: str) -> Set[QualifiedName]:
@@ -358,7 +376,7 @@ def f(): pass
             else:
                 import f
             import a.b as f
-            
+
             f()
             """
         )