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
6
7from pathlib import Path
8from typing import Collection, Dict, List, Mapping, TYPE_CHECKING
9
10import libcst as cst
11from libcst._types import StrPath
12from libcst.metadata.wrapper import MetadataWrapper
13
14if TYPE_CHECKING:
15 from libcst.metadata.base_provider import ProviderT # noqa: F401
16
17
18class FullRepoManager:
19 def __init__(
20 self,
21 repo_root_dir: StrPath,
22 paths: Collection[str],
23 providers: Collection["ProviderT"],
24 timeout: int = 5,
25 use_pyproject_toml: bool = False,
26 ) -> None:
27 """
28 Given project root directory with pyre and watchman setup, :class:`~libcst.metadata.FullRepoManager`
29 handles the inter process communication to read the required full repository cache data for
30 metadata provider like :class:`~libcst.metadata.TypeInferenceProvider`.
31
32 :param paths: a collection of paths to access full repository data.
33 :param providers: a collection of metadata provider classes require accessing full repository data, currently supports
34 :class:`~libcst.metadata.TypeInferenceProvider` and
35 :class:`~libcst.metadata.FullyQualifiedNameProvider`.
36 :param timeout: number of seconds. Raises `TimeoutExpired <https://docs.python.org/3/library/subprocess.html#subprocess.TimeoutExpired>`_
37 when timeout.
38 """
39 self.root_path: Path = Path(repo_root_dir)
40 self._cache: Dict["ProviderT", Mapping[str, object]] = {}
41 self._timeout = timeout
42 self._use_pyproject_toml = use_pyproject_toml
43 self._providers = providers
44 self._paths: List[str] = list(paths)
45
46 @property
47 def cache(self) -> Dict["ProviderT", Mapping[str, object]]:
48 """
49 The full repository cache data for all metadata providers passed in the ``providers`` parameter when
50 constructing :class:`~libcst.metadata.FullRepoManager`. Each provider is mapped to a mapping of path to cache.
51 """
52 # Make sure that the cache is available to us. If resolve_cache() was called manually then this is a noop.
53 self.resolve_cache()
54 return self._cache
55
56 def resolve_cache(self) -> None:
57 """
58 Resolve cache for all providers that require it. Normally this is called by
59 :meth:`~FullRepoManager.get_cache_for_path` so you do not need to call it
60 manually. However, if you intend to do a single cache resolution pass before
61 forking, it is a good idea to call this explicitly to control when cache
62 resolution happens.
63 """
64 if not self._cache:
65 cache: Dict["ProviderT", Mapping[str, object]] = {}
66 for provider in self._providers:
67 handler = provider.gen_cache
68 if handler:
69 cache[provider] = handler(
70 self.root_path,
71 self._paths,
72 timeout=self._timeout,
73 use_pyproject_toml=self._use_pyproject_toml,
74 )
75 self._cache = cache
76
77 def get_cache_for_path(self, path: str) -> Mapping["ProviderT", object]:
78 """
79 Retrieve cache for a source file. The file needs to appear in the ``paths`` parameter when
80 constructing :class:`~libcst.metadata.FullRepoManager`.
81
82 .. code-block:: python
83
84 manager = FullRepoManager(".", {"a.py", "b.py"}, {TypeInferenceProvider})
85 MetadataWrapper(module, cache=manager.get_cache_for_path("a.py"))
86 """
87 if path not in self._paths:
88 raise ValueError(
89 "The path needs to be in paths parameter when constructing FullRepoManager for efficient batch processing."
90 )
91 # Make sure that the cache is available to us. If the user called
92 # resolve_cache() manually then this is a noop.
93 self.resolve_cache()
94 return {
95 provider: data
96 for provider, files in self._cache.items()
97 for _path, data in files.items()
98 if _path == path
99 }
100
101 def get_metadata_wrapper_for_path(self, path: str) -> MetadataWrapper:
102 """
103 Create a :class:`~libcst.metadata.MetadataWrapper` given a source file path.
104 The path needs to be a path relative to project root directory.
105 The source code is read and parsed as :class:`~libcst.Module` for
106 :class:`~libcst.metadata.MetadataWrapper`.
107
108 .. code-block:: python
109
110 manager = FullRepoManager(".", {"a.py", "b.py"}, {TypeInferenceProvider})
111 wrapper = manager.get_metadata_wrapper_for_path("a.py")
112 """
113 module = cst.parse_module((self.root_path / path).read_text())
114 cache = self.get_cache_for_path(path)
115 return MetadataWrapper(module, True, cache)