1# Copyright (c) Meta Platforms, Inc. and affiliates.
2#
3# This source code is licensed under the MIT license found in the
4# LICENSE file in the root directory of this source tree.
5
6from pathlib import Path
7from types import MappingProxyType
8from typing import (
9 Generic,
10 List,
11 Mapping,
12 MutableMapping,
13 Optional,
14 Protocol,
15 Type,
16 TYPE_CHECKING,
17 TypeVar,
18 Union,
19)
20
21from libcst._batched_visitor import BatchableCSTVisitor
22from libcst._metadata_dependent import (
23 _T as _MetadataT,
24 _UNDEFINED_DEFAULT,
25 LazyValue,
26 MetadataDependent,
27)
28from libcst._visitors import CSTVisitor
29
30if TYPE_CHECKING:
31 from libcst._nodes.base import CSTNode
32 from libcst._nodes.module import _ModuleSelfT as _ModuleT, Module
33 from libcst.metadata.wrapper import MetadataWrapper
34
35
36ProviderT = Type["BaseMetadataProvider[object]"]
37# BaseMetadataProvider[int] would be a subtype of BaseMetadataProvider[object], so the
38# typevar is covariant.
39_ProvidedMetadataT = TypeVar("_ProvidedMetadataT", covariant=True)
40MaybeLazyMetadataT = Union[LazyValue[_ProvidedMetadataT], _ProvidedMetadataT]
41
42
43class GenCacheMethod(Protocol):
44 def __call__(
45 self,
46 root_path: Path,
47 paths: List[str],
48 *,
49 timeout: Optional[int] = None,
50 use_pyproject_toml: bool = False,
51 ) -> Mapping[str, object]: ...
52
53
54# We can't use an ABCMeta here, because of metaclass conflicts
55class BaseMetadataProvider(MetadataDependent, Generic[_ProvidedMetadataT]):
56 """
57 The low-level base class for all metadata providers. This class should be
58 extended for metadata providers that are not visitor-based.
59
60 This class is generic. A subclass of ``BaseMetadataProvider[T]`` will
61 provider metadata of type ``T``.
62 """
63
64 #: Cache of metadata computed by this provider
65 #
66 # N.B. This has some typing variance problems. See `set_metadata` for an
67 # explanation.
68 _computed: MutableMapping["CSTNode", MaybeLazyMetadataT]
69
70 #: Implement gen_cache to indicate the metadata provider depends on cache from external
71 #: system. This function will be called by :class:`~libcst.metadata.FullRepoManager`
72 #: to compute required cache object per file path.
73 gen_cache: Optional[GenCacheMethod] = None
74
75 def __init__(self, cache: object = None) -> None:
76 super().__init__()
77 self._computed: MutableMapping["CSTNode", MaybeLazyMetadataT] = {}
78 if self.gen_cache and cache is None:
79 # The metadata provider implementation is responsible to store and use cache.
80 raise ValueError(
81 f"Cache is required for initializing {self.__class__.__name__}."
82 )
83 self.cache = cache
84
85 def _gen(
86 self, wrapper: "MetadataWrapper"
87 ) -> Mapping["CSTNode", MaybeLazyMetadataT]:
88 """
89 Resolves and returns metadata mapping for the module in ``wrapper``.
90
91 This method is used by the metadata resolver and should not be called
92 directly.
93 """
94
95 self._computed = {}
96 # Resolve metadata dependencies for this provider
97 with self.resolve(wrapper):
98 self._gen_impl(wrapper.module)
99
100 # Copy into a mapping proxy to ensure immutability
101 return MappingProxyType(dict(self._computed))
102
103 def _gen_impl(self, module: "Module") -> None:
104 """
105 Override this method with a metadata computation implementation.
106 """
107 ...
108
109 def set_metadata(self, node: "CSTNode", value: MaybeLazyMetadataT) -> None:
110 """
111 Record a metadata value ``value`` for ``node``.
112 """
113 self._computed[node] = value
114
115 def get_metadata(
116 self,
117 key: Type["BaseMetadataProvider[_MetadataT]"],
118 node: "CSTNode",
119 default: Union[
120 MaybeLazyMetadataT, Type[_UNDEFINED_DEFAULT]
121 ] = _UNDEFINED_DEFAULT,
122 ) -> _MetadataT:
123 """
124 The same method as :func:`~libcst.MetadataDependent.get_metadata` except
125 metadata is accessed from ``self._computed`` in addition to ``self.metadata``.
126 See :func:`~libcst.MetadataDependent.get_metadata`.
127 """
128 if key is type(self):
129 if default is not _UNDEFINED_DEFAULT:
130 ret = self._computed.get(node, default)
131 else:
132 ret = self._computed[node]
133 if isinstance(ret, LazyValue):
134 return ret()
135 return ret
136
137 return super().get_metadata(key, node, default)
138
139
140class VisitorMetadataProvider(CSTVisitor, BaseMetadataProvider[_ProvidedMetadataT]):
141 """
142 The low-level base class for all non-batchable visitor-based metadata
143 providers. Inherits from :class:`~libcst.CSTVisitor`.
144
145 This class is generic. A subclass of ``VisitorMetadataProvider[T]`` will
146 provider metadata of type ``T``.
147 """
148
149 def _gen_impl(self, module: "_ModuleT") -> None:
150 module.visit(self)
151
152
153class BatchableMetadataProvider(
154 BatchableCSTVisitor, BaseMetadataProvider[_ProvidedMetadataT]
155):
156 """
157 The low-level base class for all batchable visitor-based metadata providers.
158 Batchable providers should be preferred when possible as they are more
159 efficient to run compared to non-batchable visitor-based providers.
160 Inherits from :class:`~libcst.BatchableCSTVisitor`.
161
162 This class is generic. A subclass of ``BatchableMetadataProvider[T]`` will
163 provider metadata of type ``T``.
164 """
165
166 def _gen_impl(self, module: "Module") -> None:
167 """
168 Batchables providers are resolved through _gen_batchable] so no
169 implementation should be provided in _gen_impl.
170 """
171 pass