From 78f4b1da05bb749b9b33143689049e99c9d4941a Mon Sep 17 00:00:00 2001 From: Roman Inflianskas Date: Tue, 6 Dec 2022 07:44:11 +0200 Subject: [PATCH] RFC and WIP: Add MypyTypeInferenceProvider This change is RFC and WIP (please read whole change message). Add `MypyTypeInferenceProvider` as an alternative for `TypeInferenceProvider`. The provider infers types using mypy as library. The only requirement for the usage is to have the latest mypy installed. Types inferred are mypy types, since mypy type system is well designed, to avoid the conversion, and also to keep it simple. For compatibility and extensibility reasons, these types are stored in separate field `MypyType.mypy_type`. Let's assume we have the following code in the file `x.py` which we want to inspect: ```python x = [42] s = set() from enum import Enum class E(Enum): f = "f" e = E.f ``` Then to get play with mypy types one should use the code like: ```python import libcst as cst from libcst.metadata import MypyTypeInferenceProvider filename = "x.py" module = cst.parse_module(open(filename).read()) cache = MypyTypeInferenceProvider.gen_cache(".", [filename])[filename] wrapper = cst.MetadataWrapper( module, cache={MypyTypeInferenceProvider: cache}, ) mypy_type = wrapper.resolve(MypyTypeInferenceProvider) x_name_node = wrapper.module.body[0].body[0].targets[0].target set_call_node = wrapper.module.body[1].body[0].value e_name_node = wrapper.module.body[-1].body[0].targets[0].target print(mypy_type[x_name_node]) # prints: builtins.list[builtins.int] print(mypy_type[x_name_node].fullname) # prints: builtins.list[builtins.int] print(mypy_type[x_name_node].mypy_type.type.fullname) # prints: builtins.list print(mypy_type[x_name_node].mypy_type.args) # prints: (builtins.int,) print(mypy_type[x_name_node].mypy_type.type.bases[0].type.fullname) # prints: typing.MutableSequence print(mypy_type[set_call_node]) # prints: builtins.set print("issuperset" in mypy_type[set_call_node].mypy_type.names) # prints: True print(mypy_type[set_call_node.func]) # prints: typing.Type[builtins.set] print(mypy_type[e_name_node].mypy_type.type.is_enum) # prints: True ``` Why? 1. `TypeInferenceProvider` requires pyre (`pyre-check` on PyPI) to be installed. mypy is more popular than pyre. If the organization uses mypy already (which is almost always the case), it may be difficult to assure colleagues (including security team) that "we need yet another type checker". 2. Even though it is possible to run pyre without watchman installation, this is not advertised. watchman installation is not always possible because of system requirements, or because of the security requirements like "we install only our favorite GNU/Linux distribution packages". 3. `TypeInferenceProvider` usage requires `pyre start` command to be run before the execution, and `pyre stop` - after the execution. This may be inconvenient, especially for the cases when pyre was not used before. 4. Types produced by pyre in `TypeInferenceProvider` are just strings. For example, it's not easily possible to infer that some variable is enum instance. `MypyTypeInferenceProvider` makes it easy, see the code above. Drawbacks: 1. Speed. mypy is slower than pyre, so is `MypyTypeInferenceProvider` comparing to `TypeInferenceProvider`. How to partially solve this: 1. Implement AST tree caching in mypy. It may be difficult, however this will lead to speed improvements for all the projects that use this functionality. 2. Implement inferred types caching inside LibCST. As far as I know, no caching at all is implemented inside LibCST, which is the prerequisite for inferred types caching, so the task is big. 3. Implement LibCST CST to mypy AST. I am not sure if this possible at all. Even if it is possible, the task is huge. 2. Two providers are doing similar things in LibCST will be present, this can potentially lead to the situation when there is a need install two type checkers to get all codemods from the library running. Alternatives considered: 1. Put `MypyTypeInferenceProvider` inside separate library (say, LibCST-mypy or `libcst-mypy` on PyPI). This will explicitly separate `MypyTypeInferenceProvider` from the rest of LibCST. Drawbacks: 1. The need to maintain separate library. 2. Limited fame (people need to know that the library exists). 3. Since some codemods cannot be implemented easily without the library, for example, `if-elif-else` to `match` converter (it needs powerful type inference), they are doomed to not be shipped with LibCST. 2. Implement base class for inferred type, which inherits from `str` (to keep the compatibility with the existing codebase) and the mechanism for dynamically selecting `TypeInferenceProvider` type checker (mypy or pyre; user can do this via environmental variable). If the code inside LibCST requires just shallow type information (so, just `str` is enough), then the code can run with any type checker. The remaining code (such as `if-elif-else` to `match` converter) will still require mypy. Misc: Code does not lint in my env, by some reason `pyre check` cannot find `mypy` library. --- libcst/metadata/__init__.py | 2 + .../metadata/mypy_type_inference_provider.py | 96 ++++++++++++ libcst/metadata/mypy_utils.py | 146 ++++++++++++++++++ .../test_mypy_type_inference_provider.py | 63 ++++++++ requirements-dev.txt | 1 + 5 files changed, 308 insertions(+) create mode 100644 libcst/metadata/mypy_type_inference_provider.py create mode 100644 libcst/metadata/mypy_utils.py create mode 100644 libcst/metadata/tests/test_mypy_type_inference_provider.py diff --git a/libcst/metadata/__init__.py b/libcst/metadata/__init__.py index 66e7e5251..ab33b2bd3 100644 --- a/libcst/metadata/__init__.py +++ b/libcst/metadata/__init__.py @@ -17,6 +17,7 @@ ExpressionContextProvider, ) from libcst.metadata.full_repo_manager import FullRepoManager +from libcst.metadata.mypy_type_inference_provider import MypyTypeInferenceProvider from libcst.metadata.name_provider import ( FullyQualifiedNameProvider, QualifiedNameProvider, @@ -74,6 +75,7 @@ "ClassScope", "ComprehensionScope", "ScopeProvider", + "MypyTypeInferenceProvider", "ParentNodeProvider", "QualifiedName", "QualifiedNameSource", diff --git a/libcst/metadata/mypy_type_inference_provider.py b/libcst/metadata/mypy_type_inference_provider.py new file mode 100644 index 000000000..29e9a6b99 --- /dev/null +++ b/libcst/metadata/mypy_type_inference_provider.py @@ -0,0 +1,96 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path +from typing import Dict, List, Mapping, Optional, TYPE_CHECKING + +import libcst as cst +from libcst._position import CodeRange +from libcst.helpers import calculate_module_and_package +from libcst.metadata.base_provider import BatchableMetadataProvider +from libcst.metadata.position_provider import PositionProvider + +try: + import mypy + + MYPY_INSTALLED = True +except ImportError: + MYPY_INSTALLED = False + + +if TYPE_CHECKING: + import mypy.nodes + + import libcst.metadata.mypy_utils + + +def raise_on_mypy_non_installed() -> None: + if not MYPY_INSTALLED: + raise RuntimeError("mypy is not installed, please install it") + + +class MypyTypeInferenceProvider( + BatchableMetadataProvider["libcst.metadata.mypy_utils.MypyType"] +): + """ + Access inferred type annotation through `mypy `_. + """ + + METADATA_DEPENDENCIES = (PositionProvider,) + + @classmethod + def gen_cache( + cls, root_path: Path, paths: List[str], timeout: Optional[int] = None + ) -> Mapping[ + str, Optional["libcst.metadata.mypy_utils.MypyTypeInferenceProviderCache"] + ]: + raise_on_mypy_non_installed() + + import mypy.build + import mypy.main + + from libcst.metadata.mypy_utils import MypyTypeInferenceProviderCache + + targets, options = mypy.main.process_options(paths) + options.preserve_asts = True + options.fine_grained_incremental = True + options.use_fine_grained_cache = True + mypy_result = mypy.build.build(targets, options=options) + cache = {} + for path in paths: + module = calculate_module_and_package(str(root_path), path).name + cache[path] = MypyTypeInferenceProviderCache( + module_name=module, + mypy_file=mypy_result.graph[module].tree, + ) + return cache + + def __init__( + self, + cache: Optional["libcst.metadata.mypy_utils.MypyTypeInferenceProviderCache"], + ) -> None: + from libcst.metadata.mypy_utils import CodeRangeToMypyNodesBinder + + super().__init__(cache) + self._mypy_node_locations: Dict[CodeRange, "mypy.nodes.Node"] = {} + if cache is None: + return + code_range_to_mypy_nodes_binder = CodeRangeToMypyNodesBinder(cache.module_name) + code_range_to_mypy_nodes_binder.visit_mypy_file(cache.mypy_file) + self._mypy_node_locations = code_range_to_mypy_nodes_binder.locations + + def _parse_metadata(self, node: cst.CSTNode) -> None: + range = self.get_metadata(PositionProvider, node) + if range in self._mypy_node_locations: + self.set_metadata(node, self._mypy_node_locations.get(range)) + + def visit_Name(self, node: cst.Name) -> Optional[bool]: + self._parse_metadata(node) + + def visit_Attribute(self, node: cst.Attribute) -> Optional[bool]: + self._parse_metadata(node) + + def visit_Call(self, node: cst.Call) -> Optional[bool]: + self._parse_metadata(node) diff --git a/libcst/metadata/mypy_utils.py b/libcst/metadata/mypy_utils.py new file mode 100644 index 000000000..17b3ce185 --- /dev/null +++ b/libcst/metadata/mypy_utils.py @@ -0,0 +1,146 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass, field +from typing import Dict, Optional, Union + +import mypy.build +import mypy.main +import mypy.modulefinder +import mypy.nodes +import mypy.options +import mypy.patterns +import mypy.traverser +import mypy.types +import mypy.typetraverser + +from libcst._add_slots import add_slots +from libcst._position import CodePosition, CodeRange + + +@add_slots +@dataclass(frozen=True) +class MypyTypeInferenceProviderCache: + module_name: str + mypy_file: mypy.nodes.MypyFile + + +@add_slots +@dataclass(frozen=True) +class MypyType: + is_type_constructor: bool + mypy_type: Union[mypy.types.Type, mypy.nodes.TypeInfo] + fullname: str = field(init=False) + + def __post_init__(self) -> None: + if isinstance(self.mypy_type, mypy.types.Type): + fullname = str(self.mypy_type) + else: + fullname = self.mypy_type.fullname + if self.is_type_constructor: + fullname = f"typing.Type[{fullname}]" + object.__setattr__(self, "fullname", fullname) + + def __str__(self) -> str: + return self.fullname + + +class CodeRangeToMypyNodesBinder( + mypy.traverser.TraverserVisitor, mypy.typetraverser.TypeTraverserVisitor +): + def __init__(self, module_name: str) -> None: + super().__init__() + self.locations: Dict[CodeRange, MypyType] = {} + self.in_type_alias_expr = False + self.module_name = module_name + + # Helpers + + @staticmethod + def get_code_range(o: mypy.nodes.Context) -> CodeRange: + return CodeRange( + start=CodePosition(o.line, o.column), + end=CodePosition(o.end_line, o.end_column), + ) + + @staticmethod + def check_bounds(o: mypy.nodes.Context) -> bool: + return ( + (o.line is not None) + and (o.line >= 1) + and (o.column is not None) + and (o.column >= 0) + and (o.end_line is not None) + and (o.end_line >= 1) + and (o.end_column is not None) + and (o.end_column >= 0) + ) + + def record_type_location_using_code_range( + self, + code_range: CodeRange, + t: Optional[Union[mypy.types.Type, mypy.nodes.TypeInfo]], + is_type_constructor: bool, + ) -> None: + if t is not None: + self.locations[code_range] = MypyType( + is_type_constructor=is_type_constructor, mypy_type=t + ) + + def record_type_location( + self, + o: mypy.nodes.Context, + t: Optional[Union[mypy.types.Type, mypy.nodes.TypeInfo]], + is_type_constructor: bool, + ) -> None: + if self.check_bounds(o): + self.record_type_location_using_code_range( + code_range=self.get_code_range(o), + t=t, + is_type_constructor=is_type_constructor, + ) + + def record_location_by_name_expr( + self, code_range: CodeRange, o: mypy.nodes.NameExpr, is_type_constructor: bool + ) -> None: + if isinstance(o.node, mypy.nodes.Var): + self.record_type_location_using_code_range( + code_range=code_range, t=o.node.type, is_type_constructor=False + ) + elif isinstance(o.node, mypy.nodes.TypeInfo): + self.record_type_location_using_code_range( + code_range=code_range, t=o.node, is_type_constructor=is_type_constructor + ) + + # Actual visitors + + def visit_var(self, o: mypy.nodes.Var) -> None: + super().visit_var(o) + self.record_type_location(o=o, t=o.type, is_type_constructor=False) + + def visit_name_expr(self, o: mypy.nodes.NameExpr) -> None: + super().visit_name_expr(o) + # Implementation in base classes is omitted, record it if it is variable or class + self.record_location_by_name_expr( + self.get_code_range(o), o, is_type_constructor=True + ) + + def visit_member_expr(self, o: mypy.nodes.MemberExpr) -> None: + super().visit_member_expr(o) + # Implementation in base classes is omitted, record it + # o.def_var should not be None after mypy run, checking here just to be sure + if o.def_var is not None: + self.record_type_location(o=o, t=o.def_var.type, is_type_constructor=False) + + def visit_call_expr(self, o: mypy.nodes.CallExpr) -> None: + super().visit_call_expr(o) + if isinstance(o.callee, mypy.nodes.NameExpr): + self.record_location_by_name_expr( + code_range=self.get_code_range(o), o=o.callee, is_type_constructor=False + ) + + def visit_instance(self, o: mypy.types.Instance) -> None: + super().visit_instance(o) + self.record_type_location(o=o, t=o, is_type_constructor=False) diff --git a/libcst/metadata/tests/test_mypy_type_inference_provider.py b/libcst/metadata/tests/test_mypy_type_inference_provider.py new file mode 100644 index 000000000..1303f35fb --- /dev/null +++ b/libcst/metadata/tests/test_mypy_type_inference_provider.py @@ -0,0 +1,63 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import libcst as cst +from libcst import MetadataWrapper +from libcst.metadata.mypy_type_inference_provider import MypyTypeInferenceProvider +from libcst.testing.utils import data_provider, UnitTest +from libcst.tests.test_pyre_integration import TEST_SUITE_PATH + + +def _test_simple_class_helper(test: UnitTest, wrapper: MetadataWrapper) -> None: + mypy_nodes = wrapper.resolve(MypyTypeInferenceProvider) + m = wrapper.module + assign = cst.ensure_type( + cst.ensure_type( + cst.ensure_type( + cst.ensure_type(m.body[1].body, cst.IndentedBlock).body[0], + cst.FunctionDef, + ).body.body[0], + cst.SimpleStatementLine, + ).body[0], + cst.AnnAssign, + ) + self_number_attr = cst.ensure_type(assign.target, cst.Attribute) + test.assertEqual(str(mypy_nodes[self_number_attr]), "builtins.int") + + # self + test.assertEqual( + str(mypy_nodes[self_number_attr.value]), "libcst.tests.pyre.simple_class.Item" + ) + collector_assign = cst.ensure_type( + cst.ensure_type(m.body[3], cst.SimpleStatementLine).body[0], cst.Assign + ) + collector = collector_assign.targets[0].target + test.assertEqual( + str(mypy_nodes[collector]), "libcst.tests.pyre.simple_class.ItemCollector" + ) + items_assign = cst.ensure_type( + cst.ensure_type(m.body[4], cst.SimpleStatementLine).body[0], cst.AnnAssign + ) + items = items_assign.target + test.assertEqual( + str(mypy_nodes[items]), "typing.Sequence[libcst.tests.pyre.simple_class.Item]" + ) + + +class MypyTypeInferenceProviderTest(UnitTest): + @data_provider( + ((TEST_SUITE_PATH / "simple_class.py", TEST_SUITE_PATH / "simple_class.json"),) + ) + def test_simple_class_types(self, source_path: Path, data_path: Path) -> None: + file = str(source_path) + repo_root = Path(__file__).parents[3] + cache = MypyTypeInferenceProvider.gen_cache(repo_root, [file]) + wrapper = MetadataWrapper( + cst.parse_module(source_path.read_text()), + cache={MypyTypeInferenceProvider: cache[file]}, + ) + _test_simple_class_helper(self, wrapper) diff --git a/requirements-dev.txt b/requirements-dev.txt index 09bcd66b3..cc802708e 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -7,6 +7,7 @@ hypothesis>=4.36.0 hypothesmith>=0.0.4 jupyter>=1.0.0 maturin>=0.8.3,<0.14 +mypy>=0.991 nbsphinx>=0.4.2 prompt-toolkit>=2.0.9 pyre-check==0.9.9; platform_system != "Windows"